"""Validation for immutable LAB E9 multirate perception 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.e9-multirate-perception-result/v1"
IDENTITY_SCHEMA = "missioncore.e9-multirate-perception-identity/v1"
REPORT_SCHEMA = "missioncore.e9-multirate-perception-report/v1"
PROFILE_SCHEMA = "missioncore.e9-multirate-perception-profile/v1"
DETECTOR_FRAME_SCHEMA = "missioncore.e9-multirate-detector-frame/v1"
SEMANTIC_FRAME_SCHEMA = "missioncore.e9-multirate-semantic-frame/v1"
MERGED_FRAME_SCHEMA = "missioncore.e9-multirate-merged-frame/v1"
MAX_JSON_BYTES = 64 * 1024 * 1024
MAX_LINE_BYTES = 2 * 1024 * 1024
MAX_SAFE_INTEGER = (1 << 53) - 1

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


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


@dataclass(frozen=True, slots=True)
class MultiratePerceptionQualificationResult:
    result_id: str
    result_root: Path
    job: CameraComputeJob
    expected_detector_frames: int
    processed_detector_frames: int
    processed_semantic_frames: int
    accepted: bool
    artifacts: tuple[MultiratePerceptionArtifact, ...]

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


def validate_multirate_perception_qualification_result(
    job_root: Path,
    result_root: Path,
) -> MultiratePerceptionQualificationResult:
    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("multirate 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"e9-multirate-perception-{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("session_id") != job.session_id
        or identity.get("source_id") != job.source_id
        or result.get("ground_truth") is not False
        or result.get("publication_scope") != "recorded-multirate-qualification-only"
        or result.get("acceptance_state") not in {"accepted", "rejected"}
    ):
        raise SessionIntegrityError("multirate identity is inconsistent")
    selection = identity.get("selection")
    configuration = identity.get("configuration")
    if not isinstance(selection, dict) or not isinstance(configuration, dict):
        raise SessionIntegrityError("multirate selection is unavailable")
    expected = selection.get("frame_count")
    source_start = selection.get("source_start_frame_index")
    source_end = selection.get("source_end_frame_index")
    timeline_start = selection.get("timeline_start_seconds")
    timeline_end = selection.get("timeline_end_seconds")
    profile = configuration.get("profile")
    if (
        not isinstance(expected, int)
        or isinstance(expected, bool)
        or expected < 2
        or not isinstance(source_start, int)
        or isinstance(source_start, bool)
        or not isinstance(source_end, int)
        or isinstance(source_end, bool)
        or source_start < 0
        or source_end - source_start + 1 != expected
        or source_end >= 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 not isinstance(profile, dict)
        or profile.get("schema_version") != PROFILE_SCHEMA
        or not isinstance(profile.get("replay"), dict)
    ):
        raise SessionIntegrityError("multirate selection is invalid")
    replay = profile["replay"]
    stride = replay.get("semantic_sample_every_frames")
    if (
        not isinstance(stride, int)
        or isinstance(stride, bool)
        or stride < 2
        or replay.get("detector_queue_capacity") not in range(1, 9)
        or replay.get("semantic_queue_capacity") not in range(1, 5)
    ):
        raise SessionIntegrityError("multirate scheduling profile is invalid")

    artifacts = _validate_artifacts(root, result.get("artifacts"))
    report = _read_object(_artifact(artifacts, "multirate-run-report").path, root)
    metrics = report.get("metrics")
    acceptance = report.get("acceptance")
    if (
        report.get("schema_version") != REPORT_SCHEMA
        or report.get("result_id") != root.name
        or report.get("identity") != identity
        or report.get("ground_truth") is not False
        or report.get("state") != result.get("acceptance_state")
        or not isinstance(metrics, dict)
        or not isinstance(acceptance, dict)
        or not isinstance(acceptance.get("accepted"), bool)
        or acceptance.get("navigation_or_safety_accepted") is not False
        or (result.get("acceptance_state") == "accepted") != acceptance["accepted"]
    ):
        raise SessionIntegrityError("multirate report is inconsistent")
    detector = metrics.get("detector")
    semantic = metrics.get("semantic")
    if not isinstance(detector, dict) or not isinstance(semantic, dict):
        raise SessionIntegrityError("multirate metrics are unavailable")
    detector_processed, detector_dropped = _validate_queue_metrics(
        detector,
        expected=expected,
        capacity=int(replay["detector_queue_capacity"]),
        expected_key="frames_expected",
    )
    semantic_scheduled = ((expected - 1) // stride) + 1
    semantic_processed, _semantic_dropped = _validate_queue_metrics(
        semantic,
        expected=semantic_scheduled,
        capacity=int(replay["semantic_queue_capacity"]),
        expected_key="frames_scheduled",
    )
    if (
        result.get("frames_processed") != detector_processed
        or detector_processed + detector_dropped != expected
        or not isinstance(semantic.get("binding_status_counts"), dict)
        or sum(semantic["binding_status_counts"].values()) != detector_processed
    ):
        raise SessionIntegrityError("multirate result accounting is inconsistent")

    detector_indices = _validate_detector_rows(
        _artifact(artifacts, "multirate-detector-frames").path,
        root,
        expected_count=detector_processed,
        selection_count=expected,
        source_start=source_start,
        timeline_start=float(timeline_start),
        timeline_end=float(timeline_end),
    )
    semantic_indices = _validate_semantic_rows(
        _artifact(artifacts, "multirate-semantic-frames").path,
        root,
        expected_count=semantic_processed,
        selection_count=expected,
        source_start=source_start,
        stride=stride,
        timeline_start=float(timeline_start),
        timeline_end=float(timeline_end),
    )
    _validate_merged_rows(
        _artifact(artifacts, "multirate-merged-frames").path,
        root,
        expected_indices=detector_indices,
        source_start=source_start,
        semantic_indices=set(semantic_indices),
    )
    return MultiratePerceptionQualificationResult(
        result_id=root.name,
        result_root=root,
        job=job,
        expected_detector_frames=expected,
        processed_detector_frames=detector_processed,
        processed_semantic_frames=semantic_processed,
        accepted=bool(acceptance["accepted"]),
        artifacts=artifacts,
    )


def _validate_queue_metrics(
    metrics: dict[str, Any],
    *,
    expected: int,
    capacity: int,
    expected_key: str,
) -> tuple[int, int]:
    processed = metrics.get("frames_processed")
    dropped = metrics.get("frames_dropped")
    queue = metrics.get("queue")
    if (
        metrics.get(expected_key) != expected
        or not isinstance(processed, int)
        or isinstance(processed, bool)
        or not 1 <= processed <= expected
        or not isinstance(dropped, int)
        or isinstance(dropped, bool)
        or dropped < 0
        or processed + dropped != expected
        or not isinstance(queue, dict)
        or queue.get("capacity") != capacity
        or queue.get("published") != expected
        or queue.get("consumed") != processed
        or queue.get("dropped_overflow") != dropped
        or queue.get("final_depth") != 0
        or queue.get("closed") is not True
        or not isinstance(queue.get("maximum_depth"), int)
        or not 0 <= queue["maximum_depth"] <= capacity
    ):
        raise SessionIntegrityError("multirate queue accounting is inconsistent")
    return processed, dropped


def _validate_detector_rows(
    path: Path,
    root: Path,
    *,
    expected_count: int,
    selection_count: int,
    source_start: int,
    timeline_start: float,
    timeline_end: float,
) -> tuple[int, ...]:
    rows = _read_jsonl(path, root)
    indices: list[int] = []
    previous_time = -math.inf
    for value in rows:
        index = value.get("frame_index")
        session_seconds = value.get("session_seconds")
        if (
            value.get("schema_version") != DETECTOR_FRAME_SCHEMA
            or not isinstance(index, int)
            or isinstance(index, bool)
            or not 0 <= index < selection_count
            or (indices and index <= indices[-1])
            or value.get("source_frame_index") != source_start + index
            or not _finite_number(session_seconds)
            or not timeline_start - 0.001 <= float(session_seconds) <= timeline_end + 0.001
            or float(session_seconds) <= previous_time
            or not _finite_number(value.get("result_age_ms"))
            or float(value["result_age_ms"]) < 0
            or not isinstance(value.get("detections"), list)
            or not isinstance(value.get("tracks"), list)
        ):
            raise SessionIntegrityError("multirate detector row is invalid")
        indices.append(index)
        previous_time = float(session_seconds)
    if len(indices) != expected_count:
        raise SessionIntegrityError("multirate detector row count changed")
    return tuple(indices)


def _validate_semantic_rows(
    path: Path,
    root: Path,
    *,
    expected_count: int,
    selection_count: int,
    source_start: int,
    stride: int,
    timeline_start: float,
    timeline_end: float,
) -> tuple[int, ...]:
    rows = _read_jsonl(path, root)
    indices: list[int] = []
    previous_time = -math.inf
    for value in rows:
        index = value.get("frame_index")
        session_seconds = value.get("session_seconds")
        class_pixels = value.get("class_pixels")
        fractions = value.get("class_fractions")
        if (
            value.get("schema_version") != SEMANTIC_FRAME_SCHEMA
            or not isinstance(index, int)
            or isinstance(index, bool)
            or not 0 <= index < selection_count
            or index % stride != 0
            or (indices and index <= indices[-1])
            or value.get("source_frame_index") != source_start + index
            or not _finite_number(session_seconds)
            or not timeline_start - 0.001 <= float(session_seconds) <= timeline_end + 0.001
            or float(session_seconds) <= previous_time
            or not _finite_number(value.get("completion_age_ms"))
            or float(value["completion_age_ms"]) < 0
            or not isinstance(value.get("mask_sha256"), str)
            or _SHA256.fullmatch(value["mask_sha256"]) is None
            or not isinstance(class_pixels, dict)
            or not isinstance(fractions, dict)
            or set(class_pixels) != set(fractions)
            or any(not isinstance(count, int) or count < 0 for count in class_pixels.values())
            or any(
                not _finite_number(fraction) or not 0 <= float(fraction) <= 1
                for fraction in fractions.values()
            )
        ):
            raise SessionIntegrityError("multirate semantic row is invalid")
        indices.append(index)
        previous_time = float(session_seconds)
    if len(indices) != expected_count:
        raise SessionIntegrityError("multirate semantic row count changed")
    return tuple(indices)


def _validate_merged_rows(
    path: Path,
    root: Path,
    *,
    expected_indices: tuple[int, ...],
    source_start: int,
    semantic_indices: set[int],
) -> None:
    rows = _read_jsonl(path, root)
    indices: list[int] = []
    for value in rows:
        index = value.get("frame_index")
        semantic = value.get("semantic")
        if (
            value.get("schema_version") != MERGED_FRAME_SCHEMA
            or not isinstance(index, int)
            or value.get("source_frame_index") != source_start + index
            or not isinstance(value.get("tracks"), list)
            or not _finite_number(value.get("detector_result_age_ms"))
            or not isinstance(semantic, dict)
            or semantic.get("status") not in {"fresh", "stale", "unavailable"}
        ):
            raise SessionIntegrityError("multirate merged row is invalid")
        semantic_source = semantic.get("source_frame_index")
        if semantic["status"] == "unavailable":
            if any(
                semantic.get(key) is not None
                for key in (
                    "source_frame_index",
                    "source_age_ms",
                    "completion_age_ms",
                    "mask_sha256",
                )
            ):
                raise SessionIntegrityError("unavailable semantic binding has invented data")
        elif (
            not isinstance(semantic_source, int)
            or semantic_source > source_start + index
            or semantic_source - source_start not in semantic_indices
            or not _finite_number(semantic.get("source_age_ms"))
            or float(semantic["source_age_ms"]) < 0
            or not _finite_number(semantic.get("completion_age_ms"))
            or not isinstance(semantic.get("mask_sha256"), str)
            or _SHA256.fullmatch(semantic["mask_sha256"]) is None
        ):
            raise SessionIntegrityError("semantic binding is inconsistent")
        indices.append(index)
    if tuple(indices) != expected_indices:
        raise SessionIntegrityError("multirate merged binding changed")


def _validate_artifacts(
    root: Path,
    raw: object,
) -> tuple[MultiratePerceptionArtifact, ...]:
    expected = {
        "multirate-detector-frames": (
            "detector-frames.jsonl",
            "application/x-ndjson",
            DETECTOR_FRAME_SCHEMA,
        ),
        "multirate-semantic-frames": (
            "semantic-frames.jsonl",
            "application/x-ndjson",
            SEMANTIC_FRAME_SCHEMA,
        ),
        "multirate-merged-frames": (
            "merged-frames.jsonl",
            "application/x-ndjson",
            MERGED_FRAME_SCHEMA,
        ),
        "worker-gpu-telemetry": ("gpu-telemetry.jsonl", "application/x-ndjson", None),
        "multirate-run-report": ("run-report.json", "application/json", REPORT_SCHEMA),
    }
    if not isinstance(raw, list) or len(raw) != len(expected):
        raise SessionIntegrityError("multirate artifact set is invalid")
    artifacts: list[MultiratePerceptionArtifact] = []
    seen: set[str] = set()
    for value in raw:
        if not isinstance(value, dict):
            raise SessionIntegrityError("multirate artifact descriptor is invalid")
        kind = value.get("kind")
        if not isinstance(kind, str) or kind in seen or kind not in expected:
            raise SessionIntegrityError("multirate 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("multirate artifact identity changed")
        artifacts.append(
            MultiratePerceptionArtifact(
                kind=kind,
                path=path,
                media_type=expected_media_type,
                byte_length=byte_length,
                sha256=digest,
                schema_version=expected_schema,
            )
        )
    return tuple(artifacts)


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


def _read_jsonl(path: Path, root: Path) -> list[dict[str, Any]]:
    _confined_regular_file(path, root)
    rows: list[dict[str, Any]] = []
    try:
        with path.open("rb") as stream:
            for line in stream:
                if not line or len(line) > MAX_LINE_BYTES:
                    raise SessionIntegrityError("multirate JSONL line is outside bounds")
                value = json.loads(line)
                if not isinstance(value, dict):
                    raise SessionIntegrityError("multirate JSONL row is invalid")
                rows.append(value)
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise SessionIntegrityError("multirate JSONL is unavailable") from exc
    return rows


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