from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Final, cast

import yaml
from yaml.constructor import ConstructorError
from yaml.tokens import AliasToken, AnchorToken, TagToken

MAX_FACTORY_YAML_BYTES: Final = 64 * 1024
CAMERA_KEYS: Final = ("camera_0", "camera_1", "camera_2", "camera_3")
MAIN_CAMERA_SLOT_BY_SOURCE: Final = {
    "sensor.camera.left": "camera_0",
    "sensor.camera.right": "camera_1",
}
MAIN_CAMERA_RTSP_PATH_BY_SOURCE: Final = {
    "sensor.camera.left": "/live/chn_left_main",
    "sensor.camera.right": "/live/chn_right_main",
}
EXPECTED_CAMERA_RESOLUTION: Final = {
    "camera_0": (4000, 3000),
    "camera_1": (4000, 3000),
    "camera_2": (1280, 800),
    "camera_3": (1280, 800),
}
ADMITTED_MAIN_STREAM_RESOLUTION: Final = (800, 600)

Matrix4 = tuple[
    tuple[float, float, float, float],
    tuple[float, float, float, float],
    tuple[float, float, float, float],
    tuple[float, float, float, float],
]


class FactoryCalibrationSchemaError(ValueError):
    """A factory calibration document is unsafe or outside the reviewed K1 schema."""


class _UniqueSafeLoader(yaml.SafeLoader):
    pass


def _construct_unique_mapping(
    loader: _UniqueSafeLoader,
    node: yaml.MappingNode,
    deep: bool = False,
) -> dict[object, object]:
    mapping: dict[object, object] = {}
    for key_node, value_node in node.value:
        key = loader.construct_object(key_node, deep=deep)
        try:
            duplicate = key in mapping
        except TypeError as exc:
            raise ConstructorError(
                "while constructing a mapping",
                node.start_mark,
                "found an unhashable key",
                key_node.start_mark,
            ) from exc
        if duplicate:
            raise ConstructorError(
                "while constructing a mapping",
                node.start_mark,
                f"found duplicate key {key!r}",
                key_node.start_mark,
            )
        mapping[key] = loader.construct_object(value_node, deep=deep)
    return mapping


_UniqueSafeLoader.add_constructor(
    yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
    _construct_unique_mapping,
)


@dataclass(frozen=True)
class CameraCalibration:
    slot: str
    camera_model: str
    camera_pose: Matrix4
    distortion: tuple[float, float, float, float]
    image_width: int
    image_height: int
    intrinsic: tuple[float, float, float, float]


@dataclass(frozen=True)
class K1FactoryCalibration:
    version: str
    cameras: tuple[
        CameraCalibration,
        CameraCalibration,
        CameraCalibration,
        CameraCalibration,
    ]
    t_camera_0_from_lidar: Matrix4

    def camera(self, slot: str) -> CameraCalibration:
        for camera in self.cameras:
            if camera.slot == slot:
                return camera
        raise KeyError(slot)

    def t_camera_from_lidar(self, slot: str) -> Matrix4:
        camera = self.camera(slot)
        return _matrix_multiply(
            _rigid_inverse(camera.camera_pose),
            self.t_camera_0_from_lidar,
        )

    def normalized_profile(self) -> dict[str, object]:
        cameras: list[dict[str, object]] = []
        for camera in self.cameras:
            cameras.append(
                {
                    "slot": camera.slot,
                    "model": camera.camera_model,
                    "native_resolution": [camera.image_width, camera.image_height],
                    "intrinsic_fx_fy_cx_cy": list(camera.intrinsic),
                    "distortion_kb4": list(camera.distortion),
                    "serialized_camera_pose": {
                        "direction": f"T_camera_0_from_{camera.slot}",
                        "row_major": _matrix_as_lists(camera.camera_pose),
                    },
                    "t_camera_from_lidar": {
                        "direction": f"T_{camera.slot}_from_lidar",
                        "row_major": _matrix_as_lists(
                            self.t_camera_from_lidar(camera.slot)
                        ),
                    },
                }
            )

        streams: dict[str, object] = {}
        target_width, target_height = ADMITTED_MAIN_STREAM_RESOLUTION
        for source_id, slot in MAIN_CAMERA_SLOT_BY_SOURCE.items():
            camera = self.camera(slot)
            scale_x = target_width / camera.image_width
            scale_y = target_height / camera.image_height
            streams[source_id] = {
                "calibration_slot": slot,
                "rtsp_path": MAIN_CAMERA_RTSP_PATH_BY_SOURCE[source_id],
                "native_resolution": [camera.image_width, camera.image_height],
                "admitted_resolution": [target_width, target_height],
                "image_transform": {
                    "kind": "firmware-configured-linear-resize",
                    "scale_x": scale_x,
                    "scale_y": scale_y,
                    "crop": None,
                    "warp": None,
                },
                "admitted_intrinsic_fx_fy_cx_cy": [
                    camera.intrinsic[0] * scale_x,
                    camera.intrinsic[1] * scale_y,
                    camera.intrinsic[2] * scale_x,
                    camera.intrinsic[3] * scale_y,
                ],
                "distortion_kb4": list(camera.distortion),
                "t_camera_from_lidar": {
                    "direction": f"T_{slot}_from_lidar",
                    "row_major": _matrix_as_lists(self.t_camera_from_lidar(slot)),
                },
            }

        return {
            "schema_version": "missioncore.k1-normalized-calibration/v1",
            "vendor_schema_version": self.version,
            "transform_notation": "T_destination_from_source",
            "matrix_storage": "row-major-homogeneous-4x4",
            "translation_unit": "meter",
            "base_transform": {
                "direction": "T_camera_0_from_lidar",
                "row_major": _matrix_as_lists(self.t_camera_0_from_lidar),
            },
            "camera_pose_interpretation": (
                "serialized T_camera_0_from_camera_N; firmware xcolor inverts it "
                "before composing T_camera_N_from_lidar"
            ),
            "cameras": cameras,
            "stream_bindings": streams,
            "mapping_proof": {
                "status": "firmware-profile-verified",
                "basis": [
                    "K1 firmware main-camera declaration order",
                    "xcolor camera_N sequential loader and main-topic order",
                    "factory native resolutions matching K1 main/secondary profiles",
                ],
            },
        }


def parse_k1_factory_calibration(
    camera_yaml: bytes,
    camera_lidar_yaml: bytes,
) -> K1FactoryCalibration:
    camera_document = _load_reviewed_yaml(camera_yaml, "camera.yaml")
    extrinsic_document = _load_reviewed_yaml(
        camera_lidar_yaml,
        "extrinsic_camera_lidar.yaml",
    )
    camera_root = _exact_mapping(
        camera_document,
        {"calibrated", "version", *CAMERA_KEYS},
        "camera.yaml",
    )
    extrinsic_root = _exact_mapping(
        extrinsic_document,
        {"calibrated", "version", "transform"},
        "extrinsic_camera_lidar.yaml",
    )
    if camera_root["calibrated"] is not True:
        raise FactoryCalibrationSchemaError("camera.yaml is not marked calibrated")
    if extrinsic_root["calibrated"] is not True:
        raise FactoryCalibrationSchemaError(
            "extrinsic_camera_lidar.yaml is not marked calibrated"
        )
    camera_version = _short_text(camera_root["version"], "camera.yaml.version")
    extrinsic_version = _short_text(
        extrinsic_root["version"],
        "extrinsic_camera_lidar.yaml.version",
    )
    if camera_version != extrinsic_version:
        raise FactoryCalibrationSchemaError("factory calibration versions do not match")

    parsed_cameras = tuple(
        _parse_camera(slot, camera_root[slot]) for slot in CAMERA_KEYS
    )
    camera_tuple = cast(
        tuple[
            CameraCalibration,
            CameraCalibration,
            CameraCalibration,
            CameraCalibration,
        ],
        parsed_cameras,
    )
    if not _matrix_close(camera_tuple[0].camera_pose, _identity_matrix(), 1e-6):
        raise FactoryCalibrationSchemaError("camera_0 pose must be the identity reference")

    base_transform = _matrix4(
        extrinsic_root["transform"],
        "extrinsic_camera_lidar.yaml.transform",
    )
    _validate_rigid_transform(
        base_transform,
        "extrinsic_camera_lidar.yaml.transform",
    )
    return K1FactoryCalibration(
        version=camera_version,
        cameras=camera_tuple,
        t_camera_0_from_lidar=base_transform,
    )


def _load_reviewed_yaml(payload: bytes, label: str) -> object:
    if len(payload) > MAX_FACTORY_YAML_BYTES:
        raise FactoryCalibrationSchemaError(f"{label} exceeds the parser byte bound")
    if b"\x00" in payload:
        raise FactoryCalibrationSchemaError(f"{label} contains NUL bytes")
    try:
        text = payload.decode("utf-8", errors="strict")
    except UnicodeDecodeError as exc:
        raise FactoryCalibrationSchemaError(f"{label} is not valid UTF-8") from exc
    try:
        for token in yaml.scan(text, Loader=_UniqueSafeLoader):
            if isinstance(token, (AliasToken, AnchorToken, TagToken)):
                raise FactoryCalibrationSchemaError(
                    f"{label} contains YAML anchors, aliases or explicit tags"
                )
        value = yaml.load(text, Loader=_UniqueSafeLoader)
    except FactoryCalibrationSchemaError:
        raise
    except yaml.YAMLError as exc:
        raise FactoryCalibrationSchemaError(f"{label} is not safe valid YAML") from exc
    return value


def _exact_mapping(value: object, keys: set[str], label: str) -> dict[str, object]:
    if type(value) is not dict:
        raise FactoryCalibrationSchemaError(f"{label} must be a mapping")
    mapping = cast(dict[object, object], value)
    if not all(type(key) is str for key in mapping):
        raise FactoryCalibrationSchemaError(f"{label} keys must be strings")
    typed = cast(dict[str, object], mapping)
    if set(typed) != keys:
        raise FactoryCalibrationSchemaError(f"{label} has an unexpected key set")
    return typed


def _parse_camera(slot: str, value: object) -> CameraCalibration:
    label = f"camera.yaml.{slot}"
    node = _exact_mapping(
        value,
        {
            "camera_model",
            "camera_pose",
            "distortion",
            "image_height",
            "image_width",
            "intrinsic",
        },
        label,
    )
    model = _short_text(node["camera_model"], f"{label}.camera_model")
    if model != "kb4":
        raise FactoryCalibrationSchemaError(f"{label} must use the reviewed kb4 model")
    width = _positive_integer(node["image_width"], f"{label}.image_width")
    height = _positive_integer(node["image_height"], f"{label}.image_height")
    if (width, height) != EXPECTED_CAMERA_RESOLUTION[slot]:
        raise FactoryCalibrationSchemaError(f"{label} resolution does not match K1 FW 3.0.2")
    intrinsic = _float4(node["intrinsic"], f"{label}.intrinsic")
    fx, fy, cx, cy = intrinsic
    if fx <= 0 or fy <= 0 or not (0 <= cx <= width) or not (0 <= cy <= height):
        raise FactoryCalibrationSchemaError(f"{label}.intrinsic is not physically admissible")
    pose = _matrix4(node["camera_pose"], f"{label}.camera_pose")
    _validate_rigid_transform(pose, f"{label}.camera_pose")
    return CameraCalibration(
        slot=slot,
        camera_model=model,
        camera_pose=pose,
        distortion=_float4(node["distortion"], f"{label}.distortion"),
        image_width=width,
        image_height=height,
        intrinsic=intrinsic,
    )


def _short_text(value: object, label: str) -> str:
    if type(value) is not str:
        raise FactoryCalibrationSchemaError(f"{label} must be a string")
    text = value
    if not text or len(text) > 64 or any(ord(character) < 32 for character in text):
        raise FactoryCalibrationSchemaError(f"{label} is invalid")
    return text


def _positive_integer(value: object, label: str) -> int:
    if type(value) is not int or value <= 0:
        raise FactoryCalibrationSchemaError(f"{label} must be a positive integer")
    return value


def _finite_number(value: object, label: str) -> float:
    if type(value) not in {int, float}:
        raise FactoryCalibrationSchemaError(f"{label} must be numeric")
    number = float(cast(int | float, value))
    if not math.isfinite(number):
        raise FactoryCalibrationSchemaError(f"{label} must be finite")
    return number


def _float4(value: object, label: str) -> tuple[float, float, float, float]:
    if type(value) is not list or len(cast(list[object], value)) != 4:
        raise FactoryCalibrationSchemaError(f"{label} must contain exactly four numbers")
    numbers = tuple(
        _finite_number(item, f"{label}[{index}]")
        for index, item in enumerate(cast(list[object], value))
    )
    return cast(tuple[float, float, float, float], numbers)


def _matrix4(value: object, label: str) -> Matrix4:
    if type(value) is not list or len(cast(list[object], value)) != 16:
        raise FactoryCalibrationSchemaError(f"{label} must contain exactly 16 numbers")
    numbers = [
        _finite_number(item, f"{label}[{index}]")
        for index, item in enumerate(cast(list[object], value))
    ]
    return cast(
        Matrix4,
        tuple(tuple(numbers[row * 4 : row * 4 + 4]) for row in range(4)),
    )


def _validate_rigid_transform(matrix: Matrix4, label: str) -> None:
    if not _matrix_close_row(matrix[3], (0.0, 0.0, 0.0, 1.0), 1e-6):
        raise FactoryCalibrationSchemaError(f"{label} has an invalid homogeneous row")
    rotation = tuple(tuple(matrix[row][column] for column in range(3)) for row in range(3))
    for left in range(3):
        for right in range(3):
            dot = sum(rotation[left][axis] * rotation[right][axis] for axis in range(3))
            expected = 1.0 if left == right else 0.0
            if abs(dot - expected) > 1e-3:
                raise FactoryCalibrationSchemaError(f"{label} rotation is not orthonormal")
    determinant = (
        rotation[0][0]
        * (rotation[1][1] * rotation[2][2] - rotation[1][2] * rotation[2][1])
        - rotation[0][1]
        * (rotation[1][0] * rotation[2][2] - rotation[1][2] * rotation[2][0])
        + rotation[0][2]
        * (rotation[1][0] * rotation[2][1] - rotation[1][1] * rotation[2][0])
    )
    if abs(determinant - 1.0) > 1e-3:
        raise FactoryCalibrationSchemaError(f"{label} rotation determinant is not +1")
    if math.sqrt(sum(matrix[row][3] ** 2 for row in range(3))) > 10.0:
        raise FactoryCalibrationSchemaError(f"{label} translation is outside the meter bound")


def _rigid_inverse(matrix: Matrix4) -> Matrix4:
    rotation_transpose = tuple(
        tuple(matrix[column][row] for column in range(3)) for row in range(3)
    )
    translation = tuple(matrix[row][3] for row in range(3))
    inverted_translation = tuple(
        -sum(rotation_transpose[row][axis] * translation[axis] for axis in range(3))
        for row in range(3)
    )
    return cast(
        Matrix4,
        tuple(
            tuple((*rotation_transpose[row], inverted_translation[row]))
            for row in range(3)
        )
        + ((0.0, 0.0, 0.0, 1.0),),
    )


def _matrix_multiply(left: Matrix4, right: Matrix4) -> Matrix4:
    return cast(
        Matrix4,
        tuple(
            tuple(
                sum(left[row][axis] * right[axis][column] for axis in range(4))
                for column in range(4)
            )
            for row in range(4)
        ),
    )


def _identity_matrix() -> Matrix4:
    return (
        (1.0, 0.0, 0.0, 0.0),
        (0.0, 1.0, 0.0, 0.0),
        (0.0, 0.0, 1.0, 0.0),
        (0.0, 0.0, 0.0, 1.0),
    )


def _matrix_close(left: Matrix4, right: Matrix4, tolerance: float) -> bool:
    return all(
        abs(left[row][column] - right[row][column]) <= tolerance
        for row in range(4)
        for column in range(4)
    )


def _matrix_close_row(
    left: tuple[float, float, float, float],
    right: tuple[float, float, float, float],
    tolerance: float,
) -> bool:
    return all(abs(left[index] - right[index]) <= tolerance for index in range(4))


def _matrix_as_lists(matrix: Matrix4) -> list[list[float]]:
    return [list(row) for row in matrix]
