"""Raw frameset endpoint, frameset_seq, full calibration and the host-align comparison, on the fake camera."""

import base64
import importlib.util
from pathlib import Path

import cv2
import numpy as np
import pytest

pytest.importorskip("pyrealsense2")

from starlette.testclient import TestClient  # noqa: E402

from src.backend.api import routes, websocket  # noqa: E402
from src.backend.main import app  # noqa: E402
from src.backend.services.camera import (CameraExtrinsics, CameraIntrinsics, Distortion,  # noqa: E402
                                         RealSenseD435)
from tests.fake_realsense import FakeAlign, FakePipeline, frame_number_of_color, open_fake_camera  # noqa: E402

SCRIPT = Path(__file__).resolve().parent.parent / "scripts" / "compare_host_align.py"


def _load_script():
    spec = importlib.util.spec_from_file_location("compare_host_align", SCRIPT)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _identity_calibration(cam):
    """Same intrinsics for both streams, identity extrinsics: host align must reproduce the depth."""
    cam.intrinsics = CameraIntrinsics(610.0, 610.0, 320.0, 240.0, 610.0, 610.0, 320.0, 240.0)
    cam.distortion_ir = cam.distortion_rgb = Distortion(0, 0, 0, 0, 0)
    cam.extrinsics = CameraExtrinsics(rotation=(1, 0, 0, 0, 1, 0, 0, 0, 1), translation=(0.0, 0.0, 0.0))
    cam.stream_sizes = {"ir": (640, 480), "rgb": (640, 480)}
    cam.distortion_models = {"ir": "brown_conrady", "rgb": "brown_conrady"}


@pytest.fixture
def client():
    cam = RealSenseD435()
    open_fake_camera(cam, FakePipeline(fps=30), FakeAlign(delay_s=0.0))
    _identity_calibration(cam)
    # Each fake camera numbers its framesets from 1: start from empty shared-encoding caches.
    routes._encoded_colors.clear()
    routes._encoded_raws.clear()
    routes.set_camera(cam)
    websocket.set_camera(cam)
    yield TestClient(app)
    cam.close()
    routes.set_camera(None)
    websocket.set_camera(None)


def _decode(b64: str, flags: int) -> np.ndarray:
    return cv2.imdecode(np.frombuffer(base64.b64decode(b64), np.uint8), flags)


def test_raw_frameset_holds_color_and_depth_of_one_frameset(client):
    data = client.get("/api/raw/latest").json()["data"]
    color = _decode(data["color_jpeg_b64"], cv2.IMREAD_COLOR)
    depth = _decode(data["depth_png_b64"], cv2.IMREAD_UNCHANGED)
    assert color.shape == tuple(data["color_shape"]) and depth.shape == tuple(data["depth_shape"])
    assert depth.dtype == np.uint16 and data["depth_units"] == "mm"
    assert frame_number_of_color(color) == int(depth[0, 0]) % 256        # one frameset
    assert {"frameset_seq", "camera_timestamp_ms", "timestamp_domain", "host_timestamp"} <= set(data)


def test_raw_frameset_waits_for_a_newer_one_and_times_out(client):
    seq = client.get("/api/raw/latest").json()["data"]["frameset_seq"]
    newer = client.get("/api/raw/latest", params={"after": seq, "timeout_ms": 1000}).json()["data"]
    assert newer["frameset_seq"] > seq
    assert client.get("/api/raw/latest", params={"after": seq + 10_000, "timeout_ms": 50}).status_code == 504


def test_raw_and_color_channel_share_the_color_jpeg(client):
    """The raw payload's color JPEG is the color channel's encoding of that frameset (encoded once)."""
    data = client.get("/api/raw/latest").json()["data"]
    shared = [enc for enc in routes._encoded_colors if enc.seq == data["frameset_seq"]]
    assert len(shared) == 1
    assert base64.b64decode(data["color_jpeg_b64"]) == shared[0].jpeg()


def test_capture_carries_the_frameset_seq_of_its_frameset(client):
    capture = client.get("/api/capture").json()["data"]
    assert capture["frameset_seq"] >= 1
    color = _decode(capture["rgb_jpeg_b64"], cv2.IMREAD_COLOR)
    depth = _decode(capture["depth_png_b64"], cv2.IMREAD_UNCHANGED)
    assert frame_number_of_color(color) == int(depth[0, 0]) % 256


def test_intrinsics_have_what_an_external_align_needs(client):
    data = client.get("/api/intrinsics").json()["data"]
    for stream in ("ir", "rgb"):
        assert {"fx", "fy", "cx", "cy", "distortion", "width", "height", "model"} <= set(data[stream])
    assert data["depth_scale_m"] == pytest.approx(0.001)
    assert len(data["extrinsics_ir_to_rgb"]["rotation"]) == 9


def test_compare_depths_counts_and_differences():
    cha = _load_script()
    pc2 = np.array([[100, 0], [300, 400]], np.uint16)
    host = np.array([[100, 200], [0, 410]], np.uint16)
    result = cha.compare_depths(pc2, host)
    assert (result["valid_both"], result["only_pc2"], result["only_host"]) == (2, 1, 1)
    assert result["abs_diff_mm"]["max"] == 10 and result["equal_fraction_of_valid_both"] == 0.5


def test_comparison_script_end_to_end_on_the_fake_camera(client):
    """With identical geometry the host align reproduces the PC2 ("fake align" = identity) depth.

    Not every pixel: rs.align splats each depth pixel's footprint and keeps the
    nearest where they overlap, which on the fake's pure-noise depth changes ~3%.
    """
    cha = _load_script()
    calibration = cha._data(client.get("/api/intrinsics"))
    with cha.HostAligner(calibration) as aligner:
        result = cha.compare_one(client, aligner, tries=20)
    assert result["valid_both"] > 0.9 * result["pixels"]
    assert result["abs_diff_mm"]["median"] == 0
    assert result["equal_fraction_of_valid_both"] > 0.9