"""Validation for immutable LAB E6 tracked LiDAR qualification results."""

from __future__ import annotations

import hashlib
import json
import math
import os
import re
import stat
from dataclasses import dataclass
from pathlib import Path
from typing import Any, TypeGuard

import numpy as np

from k1link.sessions import SessionIntegrityError

from .jobs import CameraComputeJob, validate_camera_compute_job
from .perception_epoch import (
    RecordedPerceptionEpochResult,
    validate_recorded_perception_epoch_result,
)
from .tracking_qualification import (
    TrackingQualificationResult,
    validate_tracking_qualification_result,
)

RESULT_SCHEMA = "missioncore.e6-tracked-lidar-fusion-result/v1"
IDENTITY_SCHEMA = "missioncore.e6-tracked-lidar-fusion-identity/v1"
FRAME_SCHEMA = "missioncore.e6-tracked-lidar-frame/v1"
REPORT_SCHEMA = "missioncore.e6-tracked-lidar-run-report/v1"
MAX_JSON_BYTES = 64 * 1024 * 1024
MAX_FRAME_LINE_BYTES = 4 * 1024 * 1024
MAX_POINTS_PER_FRAME = 100_000
MAX_BOXES_PER_FRAME = 10_000

_SAFE_RESULT_ID = re.compile(r"^e6-fusion-[a-f0-9]{64}$")
_SHA256 = re.compile(r"^[a-f0-9]{64}$")


@dataclass(frozen=True, slots=True)
class TrackedFusionQualificationResult:
    result_id: str
    result_root: Path
    job: CameraComputeJob
    tracking: TrackingQualificationResult
    semantic: RecordedPerceptionEpochResult
    frame_count: int
    artifacts: tuple[dict[str, Any], ...]


def validate_tracked_fusion_qualification_result(
    job_root: Path,
    tracking_root: Path,
    semantic_root: Path,
    result_root: Path,
) -> TrackedFusionQualificationResult:
    job = validate_camera_compute_job(job_root)
    tracking = validate_tracking_qualification_result(job_root, tracking_root)
    semantic = validate_recorded_perception_epoch_result(job_root, semantic_root)
    root = result_root.expanduser().resolve(strict=True)
    if not root.is_dir() or _SAFE_RESULT_ID.fullmatch(root.name) is None:
        raise SessionIntegrityError("tracked fusion result root is invalid")
    result = _read_object(root / "result.json", root)
    identity = result.get("identity")
    identity_sha256 = result.get("identity_sha256")
    if (
        result.get("schema_version") != RESULT_SCHEMA
        or result.get("result_id") != root.name
        or not isinstance(identity, dict)
        or identity.get("schema_version") != IDENTITY_SCHEMA
        or not isinstance(identity_sha256, str)
        or _SHA256.fullmatch(identity_sha256) is None
        or root.name != f"e6-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 identity.get("tracking_result_id") != tracking.result_id
        or identity.get("semantic_result_id") != semantic.result_id
        or identity.get("calibration_sha256") != semantic.calibration_sha256
        or identity.get("camera_slot") != semantic.calibration_slot
        or result.get("session_id") != job.session_id
        or result.get("source_id") != job.source_id
        or result.get("ground_truth") is not False
        or result.get("publication_scope") != "qualification-clip-only"
        or result.get("frames_processed") != tracking.frame_count
    ):
        raise SessionIntegrityError("tracked fusion identity is inconsistent")
    clip = identity.get("clip")
    if (
        not isinstance(clip, dict)
        or clip.get("source_start_frame_index") != tracking.source_start_frame_index
        or clip.get("source_end_frame_index") != tracking.source_end_frame_index
        or clip.get("clip_frame_count") != tracking.frame_count
        or clip.get("timeline_start_seconds") != tracking.timeline_start_seconds
        or clip.get("timeline_end_seconds") != tracking.timeline_end_seconds
    ):
        raise SessionIntegrityError("tracked fusion clip binding changed")
    artifacts = _validate_artifacts(root, result.get("artifacts"))
    labels = _read_labels(_artifact_path(artifacts, "tracked-lidar-box-labels", root))
    _validate_arrays(
        _artifact_path(artifacts, "tracked-lidar-arrays", root),
        frame_count=tracking.frame_count,
        label_count=len(labels),
        timeline_start=tracking.timeline_start_seconds,
        timeline_end=tracking.timeline_end_seconds,
    )
    _validate_frames(
        _artifact_path(artifacts, "tracked-lidar-frame-metadata", root),
        tracking=tracking,
    )
    report = _read_object(
        _artifact_path(artifacts, "tracked-lidar-run-report", root), root
    )
    metrics = report.get("metrics")
    if (
        report.get("schema_version") != REPORT_SCHEMA
        or report.get("result_id") != root.name
        or report.get("identity") != identity
        or not isinstance(metrics, dict)
        or metrics.get("frames_expected") != tracking.frame_count
        or metrics.get("frames_processed") != tracking.frame_count
        or metrics.get("frames_fused", 0) + metrics.get("frames_depth_unavailable", 0)
        != tracking.frame_count
        or result.get("metrics") != metrics
    ):
        raise SessionIntegrityError("tracked fusion report is inconsistent")
    rrd = _artifact_path(artifacts, "tracked-lidar-rerun", root)
    if rrd.read_bytes()[:4] != b"RRF2":
        raise SessionIntegrityError("tracked fusion Rerun stream is invalid")
    return TrackedFusionQualificationResult(
        result_id=root.name,
        result_root=root,
        job=job,
        tracking=tracking,
        semantic=semantic,
        frame_count=tracking.frame_count,
        artifacts=artifacts,
    )


def _validate_artifacts(root: Path, raw: object) -> tuple[dict[str, Any], ...]:
    expected = {
        "tracked-lidar-arrays": ("fusion.npz", "application/x-npz", None),
        "tracked-lidar-frame-metadata": (
            "fusion-frames.jsonl",
            "application/x-ndjson",
            FRAME_SCHEMA,
        ),
        "tracked-lidar-box-labels": ("box-labels.json", "application/json", None),
        "tracked-lidar-overlay-video": ("fusion-overlay.mp4", "video/mp4", None),
        "tracked-lidar-rerun": ("fusion.rrd", "application/vnd.rerun.rrd", None),
        "tracked-lidar-contact-sheet": ("contact-sheet.png", "image/png", None),
        "tracked-lidar-run-report": ("run-report.json", "application/json", REPORT_SCHEMA),
    }
    if not isinstance(raw, list) or len(raw) != len(expected):
        raise SessionIntegrityError("tracked fusion artifact set is invalid")
    artifacts: list[dict[str, Any]] = []
    seen: set[str] = set()
    for value in raw:
        if not isinstance(value, dict):
            raise SessionIntegrityError("tracked fusion artifact descriptor is invalid")
        kind = value.get("kind")
        if not isinstance(kind, str) or kind in seen or kind not in expected:
            raise SessionIntegrityError("tracked fusion artifact kind is invalid")
        seen.add(kind)
        name, media_type, schema = expected[kind]
        path = root / name
        metadata = _confined_regular_file(path, root)
        if (
            value.get("path") != name
            or value.get("media_type") != media_type
            or value.get("schema_version") != schema
            or value.get("byte_length") != metadata.st_size
            or not isinstance(value.get("sha256"), str)
            or _SHA256.fullmatch(str(value["sha256"])) is None
            or _sha256(path) != value["sha256"]
        ):
            raise SessionIntegrityError("tracked fusion artifact identity changed")
        artifacts.append(value)
    return tuple(artifacts)


def _validate_arrays(
    path: Path,
    *,
    frame_count: int,
    label_count: int,
    timeline_start: float,
    timeline_end: float,
) -> None:
    required = {
        "frame_times_ns",
        "cloud_offsets",
        "cloud_points",
        "support_offsets",
        "support_points",
        "support_colors",
        "box_offsets",
        "box_centers",
        "box_half_sizes",
        "box_quaternions",
        "box_colors",
        "box_track_ids",
        "box_distances_m",
        "box_support_counts",
    }
    try:
        with np.load(path, allow_pickle=False) as arrays:
            if set(arrays.files) != required:
                raise SessionIntegrityError("tracked fusion array set is invalid")
            times = arrays["frame_times_ns"]
            cloud_offsets = arrays["cloud_offsets"]
            cloud_points = arrays["cloud_points"]
            support_offsets = arrays["support_offsets"]
            support_points = arrays["support_points"]
            support_colors = arrays["support_colors"]
            box_offsets = arrays["box_offsets"]
            centers = arrays["box_centers"]
            half_sizes = arrays["box_half_sizes"]
            quaternions = arrays["box_quaternions"]
            colors = arrays["box_colors"]
            track_ids = arrays["box_track_ids"]
            distances = arrays["box_distances_m"]
            support_counts = arrays["box_support_counts"]
            boxes = centers.shape[0]
            if (
                times.dtype != np.int64
                or times.shape != (frame_count,)
                or cloud_offsets.dtype != np.int64
                or cloud_offsets.shape != (frame_count + 1,)
                or cloud_points.dtype != np.float32
                or cloud_points.ndim != 2
                or cloud_points.shape[1:] != (3,)
                or support_offsets.dtype != np.int64
                or support_offsets.shape != (frame_count + 1,)
                or support_points.dtype != np.float32
                or support_points.ndim != 2
                or support_points.shape[1:] != (3,)
                or support_colors.dtype != np.uint8
                or support_colors.shape != support_points.shape
                or box_offsets.dtype != np.int64
                or box_offsets.shape != (frame_count + 1,)
                or centers.dtype != np.float32
                or centers.ndim != 2
                or centers.shape[1:] != (3,)
                or half_sizes.dtype != np.float32
                or half_sizes.shape != centers.shape
                or quaternions.dtype != np.float32
                or quaternions.shape != (boxes, 4)
                or colors.dtype != np.uint8
                or colors.shape != (boxes, 4)
                or track_ids.dtype != np.int64
                or track_ids.shape != (boxes,)
                or distances.dtype != np.float32
                or distances.shape != (boxes,)
                or support_counts.dtype != np.int64
                or support_counts.shape != (boxes,)
                or label_count != boxes
            ):
                raise SessionIntegrityError("tracked fusion array shapes changed")
            _validate_offsets(cloud_offsets, cloud_points.shape[0], MAX_POINTS_PER_FRAME)
            _validate_offsets(support_offsets, support_points.shape[0], MAX_POINTS_PER_FRAME)
            _validate_offsets(box_offsets, boxes, MAX_BOXES_PER_FRAME)
            if (
                np.any(np.diff(times) <= 0)
                or times[0] < round(timeline_start * 1e9) - 1_000_000
                or times[-1] > round(timeline_end * 1e9) + 1_000_000
                or not np.isfinite(cloud_points).all()
                or not np.isfinite(support_points).all()
                or not np.isfinite(centers).all()
                or not np.isfinite(half_sizes).all()
                or not np.isfinite(quaternions).all()
                or not np.isfinite(distances).all()
                or np.any(half_sizes <= 0)
                or np.any(np.linalg.norm(quaternions, axis=1) < 1e-6)
                or np.any(track_ids < 1)
                or np.any(distances <= 0)
                or np.any(support_counts < 2)
            ):
                raise SessionIntegrityError("tracked fusion arrays are inconsistent")
    except (OSError, ValueError, KeyError) as exc:
        raise SessionIntegrityError("tracked fusion arrays are unavailable") from exc


def _validate_frames(path: Path, *, tracking: TrackingQualificationResult) -> None:
    count = 0
    previous = -math.inf
    try:
        with path.open("rb") as stream:
            for frame_index, line in enumerate(stream):
                if not line or len(line) > MAX_FRAME_LINE_BYTES:
                    raise SessionIntegrityError("tracked fusion frame line is outside bounds")
                value = json.loads(line)
                session_seconds = value.get("session_seconds") if isinstance(value, dict) else None
                objects = value.get("objects") if isinstance(value, dict) else None
                if (
                    not isinstance(value, dict)
                    or value.get("schema_version") != FRAME_SCHEMA
                    or value.get("frame_index") != frame_index
                    or value.get("source_frame_index")
                    != tracking.source_start_frame_index + frame_index
                    or not _finite(session_seconds)
                    or float(session_seconds) <= previous
                    or not isinstance(objects, list)
                    or value.get("state")
                    not in {"fused", "depth-unavailable-sync-gate"}
                    or not isinstance(value.get("accepted_cuboids"), int)
                    or value["accepted_cuboids"] < 0
                ):
                    raise SessionIntegrityError("tracked fusion frame metadata changed")
                accepted = 0
                for item in objects:
                    if not isinstance(item, dict):
                        raise SessionIntegrityError("tracked fusion object is invalid")
                    status = item.get("cuboid_status")
                    if status == "accepted-point-supported-oriented-p05-p95":
                        accepted += 1
                        if (
                            not _finite(item.get("distance_smoothed_m"))
                            or not isinstance(item.get("clustered_points"), int)
                            or item["clustered_points"] < 2
                        ):
                            raise SessionIntegrityError("accepted tracked fusion object is invalid")
                if accepted != value["accepted_cuboids"]:
                    raise SessionIntegrityError("tracked fusion frame cuboid count changed")
                previous = float(session_seconds)
                count += 1
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise SessionIntegrityError("tracked fusion frame metadata is unavailable") from exc
    if count != tracking.frame_count:
        raise SessionIntegrityError("tracked fusion frame count changed")


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("tracked fusion offsets are inconsistent")


def _read_labels(path: Path) -> list[str]:
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, list) or any(
        not isinstance(item, str) or not 1 <= len(item) <= 256 for item in value
    ):
        raise SessionIntegrityError("tracked fusion labels are invalid")
    return value


def _artifact_path(
    artifacts: tuple[dict[str, Any], ...],
    kind: str,
    root: Path,
) -> Path:
    matches = [item for item in artifacts if item["kind"] == kind]
    if len(matches) != 1:
        raise SessionIntegrityError(f"tracked fusion artifact is unavailable: {kind}")
    return root / str(matches[0]["path"])


def _read_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("tracked 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("tracked fusion JSON is unavailable") from exc
    if not isinstance(value, dict):
        raise SessionIntegrityError("tracked fusion JSON root is invalid")
    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("tracked 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("tracked 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("tracked fusion identity cannot be encoded") from exc


def _finite(value: object) -> TypeGuard[int | float]:
    return (
        isinstance(value, (int, float))
        and not isinstance(value, bool)
        and math.isfinite(float(value))
    )
