"""Validated calibrated 3D fusion results projected into an opened Rerun recording."""

from __future__ import annotations

import hashlib
import json
import os
import re
import stat
import threading
from contextlib import suppress
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol

import numpy as np
import rerun as rr
from rerun.components import FillMode

from k1link.artifacts import write_json_atomic
from k1link.sessions import SessionIntegrityError

from .jobs import CameraComputeJob, validate_camera_compute_job
from .perception_epoch import validate_recorded_perception_epoch_result
from .results import RecordedPerceptionOverlayError

FUSION_SCHEMA = "missioncore.recorded-calibrated-fusion/v1"
FUSION_IDENTITY_SCHEMA = "missioncore.recorded-calibrated-fusion-identity/v1"
SESSION_TIMELINE = "session_time"
MAX_JSON_BYTES = 64 * 1024 * 1024
MAX_SCAN = 512
MAX_POINTS_PER_FRAME = 100_000
MAX_BOXES_PER_FRAME = 10_000

_SAFE_FUSION_ID = re.compile(r"^fusion-[a-f0-9]{64}$")
_SAFE_RECORDING_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$")
_SHA256 = re.compile(r"^[a-f0-9]{64}$")


@dataclass(frozen=True, slots=True)
class RecordedCalibratedFusion:
    fusion_id: str
    root: Path
    job: CameraComputeJob
    perception_result_id: str
    created_at_utc: str
    arrays_path: Path
    labels_path: Path
    frame_count: int


class _OverlayProvider(Protocol):
    def render(
        self,
        session_id: str,
        *,
        application_id: str,
        recording_id: str,
    ) -> bytes | None: ...


class RecordedCalibratedFusionStore:
    """Discover full-epoch fusion and serialize it for the opened recording ID."""

    def __init__(
        self,
        *,
        jobs_root: Path,
        perception_results_root: Path,
        fusion_results_root: Path,
        cache_root: Path,
    ) -> None:
        self.jobs_root = jobs_root.expanduser().absolute()
        self.perception_results_root = perception_results_root.expanduser().absolute()
        self.fusion_results_root = fusion_results_root.expanduser().absolute()
        self.cache_root = cache_root.expanduser().absolute()
        self._lock = threading.Lock()

    def render(
        self,
        session_id: str,
        *,
        application_id: str,
        recording_id: str,
    ) -> bytes | None:
        if _SAFE_RECORDING_ID.fullmatch(session_id) is None:
            raise ValueError("observation session id is invalid")
        if application_id != "nodedc_mission_core_recorded":
            raise ValueError("recorded fusion application id is invalid")
        if _SAFE_RECORDING_ID.fullmatch(recording_id) is None:
            raise ValueError("recorded fusion recording id is invalid")
        with self._lock:
            fusion = self._latest(session_id)
            if fusion is None:
                return None
            cache_root = _private_directory(self.cache_root)
            session_cache = _private_child_directory(cache_root, session_id)
            fusion_cache = _private_child_directory(session_cache, fusion.fusion_id)
            output = fusion_cache / f"{recording_id}.rrd"
            sidecar = output.with_suffix(".rrd.cache.json")
            cached = _read_cache(output, sidecar, fusion, recording_id)
            if cached is not None:
                return cached
            payload = _render_fusion(
                fusion,
                application_id=application_id,
                recording_id=recording_id,
            )
            temporary = output.with_name(f".{output.name}.{os.getpid()}.tmp")
            try:
                with temporary.open("xb") as stream:
                    stream.write(payload)
                    stream.flush()
                    os.fsync(stream.fileno())
                os.chmod(temporary, 0o600)
                os.replace(temporary, output)
                write_json_atomic(
                    sidecar,
                    {
                        "schema_version": "missioncore.calibrated-fusion-overlay-cache/v1",
                        "fusion_id": fusion.fusion_id,
                        "recording_id": recording_id,
                        "byte_length": len(payload),
                        "sha256": hashlib.sha256(payload).hexdigest(),
                    },
                )
            finally:
                temporary.unlink(missing_ok=True)
            return payload

    def _latest(self, session_id: str) -> RecordedCalibratedFusion | None:
        try:
            jobs = sorted(self.jobs_root.iterdir())
        except FileNotFoundError:
            return None
        if len(jobs) > MAX_SCAN:
            raise RecordedPerceptionOverlayError("compute job catalog is outside bounds")
        matches: list[RecordedCalibratedFusion] = []
        for job_root in jobs:
            if job_root.is_symlink():
                continue
            try:
                job = validate_camera_compute_job(job_root)
            except (OSError, SessionIntegrityError):
                continue
            if job.session_id != session_id:
                continue
            parent = self.fusion_results_root / job.job_id
            try:
                candidates = sorted(
                    path
                    for path in parent.iterdir()
                    if path.is_dir()
                    and not path.is_symlink()
                    and _SAFE_FUSION_ID.fullmatch(path.name) is not None
                )
            except FileNotFoundError:
                continue
            if len(candidates) > MAX_SCAN:
                raise RecordedPerceptionOverlayError("fusion result catalog is outside bounds")
            for candidate in candidates:
                try:
                    matches.append(
                        validate_recorded_calibrated_fusion(
                            job_root,
                            self.perception_results_root,
                            candidate,
                        )
                    )
                except (OSError, SessionIntegrityError):
                    continue
        if not matches:
            return None
        return max(matches, key=lambda value: (value.created_at_utc, value.fusion_id))


class RecordedPerceptionOverlayMux:
    """Prefer calibrated full-epoch 3D fusion and retain the legacy fallback."""

    def __init__(self, primary: _OverlayProvider, fallback: _OverlayProvider | None) -> None:
        self.primary = primary
        self.fallback = fallback

    def render(
        self,
        session_id: str,
        *,
        application_id: str,
        recording_id: str,
    ) -> bytes | None:
        payload = self.primary.render(
            session_id,
            application_id=application_id,
            recording_id=recording_id,
        )
        if payload is not None or self.fallback is None:
            return payload
        return self.fallback.render(
            session_id,
            application_id=application_id,
            recording_id=recording_id,
        )


def validate_recorded_calibrated_fusion(
    job_root: Path,
    perception_results_root: Path,
    fusion_root: Path,
) -> RecordedCalibratedFusion:
    job = validate_camera_compute_job(job_root)
    root = fusion_root.expanduser().resolve(strict=True)
    if not root.is_dir() or _SAFE_FUSION_ID.fullmatch(root.name) is None:
        raise SessionIntegrityError("calibrated fusion root is invalid")
    manifest = _read_json_object(root / "manifest.json", root)
    identity = manifest.get("identity")
    identity_sha256 = manifest.get("identity_sha256")
    if (
        manifest.get("schema_version") != FUSION_SCHEMA
        or manifest.get("fusion_id") != root.name
        or not isinstance(identity, dict)
        or identity.get("schema_version") != FUSION_IDENTITY_SCHEMA
        or not isinstance(identity_sha256, str)
        or _SHA256.fullmatch(identity_sha256) is None
        or root.name != f"fusion-{identity_sha256}"
        or hashlib.sha256(_canonical_json(identity)).hexdigest() != identity_sha256
        or identity.get("job_id") != job.job_id
        or identity.get("input_sha256") != job.input_sha256
        or manifest.get("session_id") != job.session_id
        or manifest.get("source_id") != job.source_id
        or manifest.get("frame_count") != job.segment_count
        or manifest.get("timeline_start_seconds") != job.timeline_start_seconds
        or manifest.get("timeline_end_seconds") != job.timeline_end_seconds
    ):
        raise SessionIntegrityError("calibrated fusion identity is inconsistent")
    perception_result_id = identity.get("perception_result_id")
    if not isinstance(perception_result_id, str):
        raise SessionIntegrityError("calibrated fusion perception binding is invalid")
    perception = validate_recorded_perception_epoch_result(
        job_root,
        perception_results_root / job.job_id / perception_result_id,
    )
    if (
        identity.get("calibration_sha256") != perception.calibration_sha256
        or identity.get("camera_slot") != perception.calibration_slot
    ):
        raise SessionIntegrityError("calibrated fusion calibration binding changed")
    artifacts = manifest.get("artifacts")
    expected = {"fusion.npz", "box-labels.json", "fusion-frames.jsonl"}
    if not isinstance(artifacts, list) or len(artifacts) != len(expected):
        raise SessionIntegrityError("calibrated fusion artifacts are incomplete")
    paths: dict[str, Path] = {}
    for artifact in artifacts:
        if not isinstance(artifact, dict) or artifact.get("name") not in expected:
            raise SessionIntegrityError("calibrated fusion artifact is invalid")
        name = str(artifact["name"])
        if name in paths:
            raise SessionIntegrityError("calibrated fusion artifact is duplicated")
        path = root / name
        metadata = _confined_regular_file(path, root)
        if (
            artifact.get("byte_length") != metadata.st_size
            or not isinstance(artifact.get("sha256"), str)
            or _SHA256.fullmatch(str(artifact["sha256"])) is None
            or _sha256(path) != artifact["sha256"]
        ):
            raise SessionIntegrityError("calibrated fusion artifact identity changed")
        paths[name] = path
    labels = _read_labels(paths["box-labels.json"])
    _validate_arrays(paths["fusion.npz"], job, len(labels))
    created_at = manifest.get("created_at_utc")
    if not isinstance(created_at, str) or not 1 <= len(created_at) <= 64:
        raise SessionIntegrityError("calibrated fusion creation time is invalid")
    return RecordedCalibratedFusion(
        fusion_id=root.name,
        root=root,
        job=job,
        perception_result_id=perception_result_id,
        created_at_utc=created_at,
        arrays_path=paths["fusion.npz"],
        labels_path=paths["box-labels.json"],
        frame_count=job.segment_count,
    )


def _validate_arrays(path: Path, job: CameraComputeJob, label_count: int) -> None:
    try:
        with np.load(path, allow_pickle=False) as arrays:
            required = {
                "frame_times_ns",
                "point_offsets",
                "points",
                "point_colors",
                "box_offsets",
                "box_centers",
                "box_half_sizes",
                "box_colors",
            }
            if set(arrays.files) != required:
                raise SessionIntegrityError("calibrated fusion array set is invalid")
            frame_times = arrays["frame_times_ns"]
            point_offsets = arrays["point_offsets"]
            points = arrays["points"]
            point_colors = arrays["point_colors"]
            box_offsets = arrays["box_offsets"]
            box_centers = arrays["box_centers"]
            box_half_sizes = arrays["box_half_sizes"]
            box_colors = arrays["box_colors"]
            if (
                frame_times.dtype != np.int64
                or frame_times.shape != (job.segment_count,)
                or point_offsets.dtype != np.int64
                or point_offsets.shape != (job.segment_count + 1,)
                or box_offsets.dtype != np.int64
                or box_offsets.shape != (job.segment_count + 1,)
                or points.dtype != np.float32
                or points.ndim != 2
                or points.shape[1:] != (3,)
                or point_colors.dtype != np.uint8
                or point_colors.shape != points.shape
                or box_centers.dtype != np.float32
                or box_centers.ndim != 2
                or box_centers.shape[1:] != (3,)
                or box_half_sizes.dtype != np.float32
                or box_half_sizes.shape != box_centers.shape
                or box_colors.dtype != np.uint8
                or box_colors.shape != (box_centers.shape[0], 4)
                or label_count != box_centers.shape[0]
            ):
                raise SessionIntegrityError("calibrated fusion array shapes changed")
            _validate_offsets(point_offsets, points.shape[0], MAX_POINTS_PER_FRAME)
            _validate_offsets(box_offsets, box_centers.shape[0], MAX_BOXES_PER_FRAME)
            if (
                np.any(np.diff(frame_times) <= 0)
                or frame_times[0] < round(job.timeline_start_seconds * 1e9) - 1_000_000
                or frame_times[-1] > round(job.timeline_end_seconds * 1e9) + 1_000_000
                or not np.isfinite(points).all()
                or not np.isfinite(box_centers).all()
                or not np.isfinite(box_half_sizes).all()
                or np.any(box_half_sizes <= 0)
            ):
                raise SessionIntegrityError("calibrated fusion arrays are inconsistent")
    except (OSError, ValueError, KeyError) as exc:
        raise SessionIntegrityError("calibrated fusion arrays are unavailable") from exc


def _validate_offsets(
    offsets: np.ndarray[Any, np.dtype[np.int64]],
    total: int,
    maximum: int,
) -> None:
    if (
        offsets[0] != 0
        or offsets[-1] != total
        or np.any(np.diff(offsets) < 0)
        or np.any(np.diff(offsets) > maximum)
    ):
        raise SessionIntegrityError("calibrated fusion offsets are inconsistent")


def _render_fusion(
    fusion: RecordedCalibratedFusion,
    *,
    application_id: str,
    recording_id: str,
) -> bytes:
    labels = _read_labels(fusion.labels_path)
    recording = rr.RecordingStream(application_id, recording_id=recording_id)
    stream = rr.binary_stream(recording)
    try:
        recording.log("/world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True)
        recording.log(
            "/world/perception/contract",
            rr.TextDocument(
                "Factory-calibrated KB4 mask-to-LiDAR diagnostic. Distances and support-gated "
                "boxes are not ground-truthed or safety accepted."
            ),
            static=True,
        )
        with np.load(fusion.arrays_path, allow_pickle=False) as arrays:
            frame_times = arrays["frame_times_ns"]
            point_offsets = arrays["point_offsets"]
            points = arrays["points"]
            point_colors = arrays["point_colors"]
            box_offsets = arrays["box_offsets"]
            centers = arrays["box_centers"]
            half_sizes = arrays["box_half_sizes"]
            box_colors = arrays["box_colors"]
            for index, timestamp in enumerate(frame_times):
                recording.set_time(
                    SESSION_TIMELINE,
                    duration=np.timedelta64(int(timestamp), "ns"),
                )
                point_start, point_end = int(point_offsets[index]), int(point_offsets[index + 1])
                if point_end > point_start:
                    recording.log(
                        "/world/perception/semantic_points",
                        rr.Points3D(
                            points[point_start:point_end],
                            colors=point_colors[point_start:point_end],
                            radii=rr.Radius.ui_points(2.5),
                        ),
                    )
                else:
                    recording.log(
                        "/world/perception/semantic_points",
                        rr.Clear(recursive=False),
                    )
                box_start, box_end = int(box_offsets[index]), int(box_offsets[index + 1])
                if box_end > box_start:
                    recording.log(
                        "/world/perception/boxes3d",
                        rr.Boxes3D(
                            centers=centers[box_start:box_end],
                            half_sizes=half_sizes[box_start:box_end],
                            colors=box_colors[box_start:box_end],
                            labels=labels[box_start:box_end],
                            fill_mode=FillMode.Solid,
                            show_labels=True,
                        ),
                    )
                else:
                    recording.log(
                        "/world/perception/boxes3d",
                        rr.Clear(recursive=False),
                    )
        payload = stream.read(flush=True, flush_timeout_sec=300.0)
    except Exception as exc:
        raise RecordedPerceptionOverlayError("failed to serialize calibrated fusion") from exc
    finally:
        with suppress(Exception):
            recording.disconnect()
    if payload is None or not payload.startswith(b"RRF2"):
        raise RecordedPerceptionOverlayError("serialized calibrated fusion is invalid")
    return payload


def _read_cache(
    output: Path,
    sidecar: Path,
    fusion: RecordedCalibratedFusion,
    recording_id: str,
) -> bytes | None:
    try:
        value = _read_json_object(sidecar, sidecar.parent)
        payload = output.read_bytes()
    except (OSError, SessionIntegrityError):
        return None
    if (
        value.get("schema_version") != "missioncore.calibrated-fusion-overlay-cache/v1"
        or value.get("fusion_id") != fusion.fusion_id
        or value.get("recording_id") != recording_id
        or value.get("byte_length") != len(payload)
        or value.get("sha256") != hashlib.sha256(payload).hexdigest()
        or not payload.startswith(b"RRF2")
    ):
        return None
    return payload


def _read_labels(path: Path) -> list[str]:
    metadata = _confined_regular_file(path, path.parent)
    if not 1 <= metadata.st_size <= MAX_JSON_BYTES:
        raise SessionIntegrityError("calibrated fusion labels are outside bounds")
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise SessionIntegrityError("calibrated fusion labels are unavailable") from exc
    if (
        not isinstance(value, list)
        or any(not isinstance(item, str) or not 1 <= len(item) <= 256 for item in value)
    ):
        raise SessionIntegrityError("calibrated fusion labels are invalid")
    return value


def _read_json_object(path: Path, root: Path) -> dict[str, Any]:
    metadata = _confined_regular_file(path, root)
    if not 0 < metadata.st_size <= MAX_JSON_BYTES:
        raise SessionIntegrityError("calibrated fusion JSON is outside bounds")
    try:
        value = json.loads(path.read_text(encoding="utf-8-sig"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise SessionIntegrityError("calibrated fusion JSON is unavailable") from exc
    if not isinstance(value, dict):
        raise SessionIntegrityError("calibrated fusion JSON is not an object")
    return value


def _confined_regular_file(path: Path, root: Path) -> os.stat_result:
    try:
        resolved_root = root.resolve(strict=True)
        resolved = path.resolve(strict=True)
        metadata = path.lstat()
    except OSError as exc:
        raise SessionIntegrityError("calibrated fusion artifact is unavailable") from exc
    if (
        stat.S_ISLNK(metadata.st_mode)
        or not stat.S_ISREG(metadata.st_mode)
        or not resolved.is_relative_to(resolved_root)
    ):
        raise SessionIntegrityError("calibrated fusion artifact is not confined")
    return metadata


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        while chunk := stream.read(1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()


def _canonical_json(value: object) -> bytes:
    try:
        return json.dumps(
            value,
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
            allow_nan=False,
        ).encode()
    except (TypeError, ValueError) as exc:
        raise SessionIntegrityError("calibrated fusion identity cannot be encoded") from exc


def _private_directory(path: Path) -> Path:
    path.mkdir(mode=0o700, parents=True, exist_ok=True)
    if path.is_symlink() or not path.is_dir():
        raise RecordedPerceptionOverlayError("fusion cache root is invalid")
    with suppress(OSError):
        os.chmod(path, 0o700)
    return path.resolve(strict=True)


def _private_child_directory(root: Path, name: str) -> Path:
    child = root / name
    child.mkdir(mode=0o700, exist_ok=True)
    if child.is_symlink() or not child.is_dir() or child.resolve(strict=True).parent != root:
        raise RecordedPerceptionOverlayError("fusion cache directory is invalid")
    with suppress(OSError):
        os.chmod(child, 0o700)
    return child.resolve(strict=True)
