"""Validation for immutable LAB E5 tracking-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

from k1link.sessions import SessionIntegrityError

from .jobs import CameraComputeJob, validate_camera_compute_job

RESULT_SCHEMA = "missioncore.e5-tracking-result/v1"
IDENTITY_SCHEMA = "missioncore.e5-tracking-identity/v1"
REPORT_SCHEMA = "missioncore.e5-tracking-report/v1"
FRAME_SCHEMA = "missioncore.e5-tracking-frame/v1"
MAX_JSON_BYTES = 64 * 1024 * 1024
MAX_FRAME_LINE_BYTES = 1024 * 1024
MAX_SAFE_INTEGER = (1 << 53) - 1

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


@dataclass(frozen=True, slots=True)
class TrackingQualificationArtifact:
    kind: str
    path: Path
    media_type: str
    byte_length: int
    sha256: str
    schema_version: str | None


@dataclass(frozen=True, slots=True)
class TrackingQualificationResult:
    result_id: str
    result_root: Path
    job: CameraComputeJob
    frame_count: int
    source_start_frame_index: int
    source_end_frame_index: int
    timeline_start_seconds: float
    timeline_end_seconds: float
    artifacts: tuple[TrackingQualificationArtifact, ...]

    def artifact(self, kind: str) -> TrackingQualificationArtifact:
        matches = tuple(item for item in self.artifacts if item.kind == kind)
        if len(matches) != 1:
            raise SessionIntegrityError(f"tracking artifact is unavailable: {kind}")
        return matches[0]


def validate_tracking_qualification_result(
    job_root: Path,
    result_root: Path,
) -> TrackingQualificationResult:
    job = validate_camera_compute_job(job_root)
    root = result_root.expanduser().resolve(strict=True)
    if not root.is_dir() or _SAFE_RESULT_ID.fullmatch(root.name) is None:
        raise SessionIntegrityError("tracking 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"e5-tracking-{identity_sha256}"
        or hashlib.sha256(_canonical_json(identity)).hexdigest() != identity_sha256
        or result.get("job_id") != job.job_id
        or result.get("input_sha256") != job.input_sha256
        or result.get("session_id") != job.session_id
        or result.get("source_id") != job.source_id
        or result.get("timestamp_basis") != "session-time-seconds"
        or result.get("ground_truth") is not False
        or result.get("publication_scope") != "qualification-clip-only"
        or identity.get("job_id") != job.job_id
        or identity.get("input_sha256") != job.input_sha256
    ):
        raise SessionIntegrityError("tracking result identity is inconsistent")
    clip = identity.get("clip")
    if not isinstance(clip, dict):
        raise SessionIntegrityError("tracking clip identity is unavailable")
    frame_count = clip.get("clip_frame_count")
    start_index = clip.get("source_start_frame_index")
    end_index = clip.get("source_end_frame_index")
    timeline_start = clip.get("timeline_start_seconds")
    timeline_end = clip.get("timeline_end_seconds")
    if (
        not isinstance(frame_count, int)
        or isinstance(frame_count, bool)
        or frame_count < 2
        or not isinstance(start_index, int)
        or isinstance(start_index, bool)
        or not isinstance(end_index, int)
        or isinstance(end_index, bool)
        or start_index < 0
        or end_index - start_index + 1 != frame_count
        or end_index >= job.segment_count
        or not _finite_number(timeline_start)
        or not _finite_number(timeline_end)
        or float(timeline_start) < job.timeline_start_seconds
        or float(timeline_end) > job.timeline_end_seconds
        or float(timeline_end) <= float(timeline_start)
        or result.get("frames_processed") != frame_count
    ):
        raise SessionIntegrityError("tracking clip identity is invalid")
    artifacts = _validate_artifacts(root, result.get("artifacts"))
    report = _read_object(_artifact(artifacts, "tracking-run-report").path, root)
    metrics = report.get("metrics")
    report_input = report.get("input")
    if (
        report.get("schema_version") != REPORT_SCHEMA
        or report.get("state") != "published-qualification"
        or report.get("result_id") != root.name
        or report.get("ground_truth") is not False
        or not isinstance(metrics, dict)
        or metrics.get("frames_expected") != frame_count
        or metrics.get("frames_processed") != frame_count
        or metrics.get("frames_failed") != 0
        or metrics.get("frames_skipped") != 0
        or not isinstance(report_input, dict)
        or report_input.get("job_id") != job.job_id
        or report_input.get("input_sha256") != job.input_sha256
        or report_input.get("source_start_frame_index") != start_index
        or report_input.get("source_end_frame_index") != end_index
    ):
        raise SessionIntegrityError("tracking report is inconsistent")
    _validate_frames(
        _artifact(artifacts, "tracking-frame-metadata"),
        frame_count=frame_count,
        source_start=start_index,
        timeline_start=float(timeline_start),
        timeline_end=float(timeline_end),
    )
    return TrackingQualificationResult(
        result_id=root.name,
        result_root=root,
        job=job,
        frame_count=frame_count,
        source_start_frame_index=start_index,
        source_end_frame_index=end_index,
        timeline_start_seconds=float(timeline_start),
        timeline_end_seconds=float(timeline_end),
        artifacts=artifacts,
    )


def _validate_artifacts(
    root: Path,
    raw: object,
) -> tuple[TrackingQualificationArtifact, ...]:
    expected = {
        "tracking-overlay-video": ("tracking.mp4", "video/mp4", None),
        "tracking-frame-metadata": ("frames.jsonl", "application/x-ndjson", FRAME_SCHEMA),
        "worker-gpu-telemetry": ("gpu-telemetry.jsonl", "application/x-ndjson", None),
        "tracking-contact-sheet": ("contact-sheet.png", "image/png", None),
        "tracking-run-report": ("run-report.json", "application/json", REPORT_SCHEMA),
    }
    if not isinstance(raw, list) or len(raw) != len(expected):
        raise SessionIntegrityError("tracking artifact set is invalid")
    artifacts: list[TrackingQualificationArtifact] = []
    seen: set[str] = set()
    for value in raw:
        if not isinstance(value, dict):
            raise SessionIntegrityError("tracking artifact descriptor is invalid")
        kind = value.get("kind")
        if not isinstance(kind, str) or kind in seen or kind not in expected:
            raise SessionIntegrityError("tracking artifact kind is invalid")
        seen.add(kind)
        expected_path, expected_media_type, expected_schema = expected[kind]
        path = root / expected_path
        metadata = _confined_regular_file(path, root)
        byte_length = value.get("byte_length")
        digest = value.get("sha256")
        if (
            value.get("path") != expected_path
            or value.get("media_type") != expected_media_type
            or value.get("schema_version") != expected_schema
            or not isinstance(byte_length, int)
            or isinstance(byte_length, bool)
            or not 0 < byte_length <= MAX_SAFE_INTEGER
            or metadata.st_size != byte_length
            or not isinstance(digest, str)
            or _SHA256.fullmatch(digest) is None
            or _sha256_file(path) != digest
        ):
            raise SessionIntegrityError("tracking artifact identity changed")
        artifacts.append(
            TrackingQualificationArtifact(
                kind=kind,
                path=path,
                media_type=expected_media_type,
                byte_length=byte_length,
                sha256=digest,
                schema_version=expected_schema,
            )
        )
    return tuple(artifacts)


def _validate_frames(
    artifact: TrackingQualificationArtifact,
    *,
    frame_count: int,
    source_start: int,
    timeline_start: float,
    timeline_end: float,
) -> None:
    previous = -math.inf
    count = 0
    try:
        with artifact.path.open("rb") as stream:
            for frame_index, line in enumerate(stream):
                if not line or len(line) > MAX_FRAME_LINE_BYTES:
                    raise SessionIntegrityError("tracking frame line is outside bounds")
                value = json.loads(line)
                session_seconds = value.get("session_seconds") if isinstance(value, dict) else None
                detections = value.get("detections") if isinstance(value, dict) else None
                tracks = value.get("tracks") 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("sequence") != frame_index + 1
                    or value.get("source_frame_index") != source_start + frame_index
                    or value.get("source_sequence") != source_start + frame_index + 1
                    or not _finite_number(session_seconds)
                    or float(session_seconds) <= previous
                    or float(session_seconds) < timeline_start - 0.001
                    or float(session_seconds) > timeline_end + 0.001
                    or not isinstance(detections, list)
                    or not isinstance(tracks, list)
                ):
                    raise SessionIntegrityError("tracking frame metadata is inconsistent")
                _validate_observations(detections, tracks)
                previous = float(session_seconds)
                count += 1
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise SessionIntegrityError("tracking frame metadata is unavailable") from exc
    if count != frame_count:
        raise SessionIntegrityError("tracking frame metadata count changed")


def _validate_observations(detections: list[object], tracks: list[object]) -> None:
    for detection in detections:
        if not isinstance(detection, dict):
            raise SessionIntegrityError("tracking detection is invalid")
        _validate_box(detection.get("bbox_xyxy"))
        if (
            not isinstance(detection.get("class_id"), int)
            or not isinstance(detection.get("label"), str)
            or not _finite_number(detection.get("score"))
            or not 0 <= float(detection["score"]) <= 1
        ):
            raise SessionIntegrityError("tracking detection fields are invalid")
    track_ids: set[int] = set()
    for track in tracks:
        if not isinstance(track, dict):
            raise SessionIntegrityError("tracking observation is invalid")
        _validate_box(track.get("bbox_xyxy"))
        track_id = track.get("track_id")
        if (
            not isinstance(track_id, int)
            or isinstance(track_id, bool)
            or track_id < 1
            or track_id in track_ids
            or not isinstance(track.get("class_id"), int)
            or not isinstance(track.get("label"), str)
            or not _finite_number(track.get("score"))
            or not isinstance(track.get("hits"), int)
            or track["hits"] < 1
        ):
            raise SessionIntegrityError("tracking observation fields are invalid")
        track_ids.add(track_id)


def _validate_box(value: object) -> None:
    if (
        not isinstance(value, list)
        or len(value) != 4
        or not all(_finite_number(item) for item in value)
        or not 0 <= float(value[0]) < float(value[2]) <= 800
        or not 0 <= float(value[1]) < float(value[3]) <= 600
    ):
        raise SessionIntegrityError("tracking box is invalid")


def _artifact(
    artifacts: tuple[TrackingQualificationArtifact, ...],
    kind: str,
) -> TrackingQualificationArtifact:
    matches = tuple(item for item in artifacts if item.kind == kind)
    if len(matches) != 1:
        raise SessionIntegrityError(f"tracking artifact is unavailable: {kind}")
    return matches[0]


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("tracking JSON is outside bounds")
    try:
        value = json.loads(path.read_text(encoding="utf-8-sig"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise SessionIntegrityError("tracking JSON is unavailable") from exc
    if not isinstance(value, dict):
        raise SessionIntegrityError("tracking 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("tracking 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("tracking artifact is not confined")
    return metadata


def _sha256_file(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("utf-8")
    except (TypeError, ValueError) as exc:
        raise SessionIntegrityError("tracking identity is not canonical JSON") from exc


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