from __future__ import annotations

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

import numpy as np
import rerun as rr

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

from .jobs import CameraComputeJob, validate_camera_compute_job

COMPUTE_RESULT_SCHEMA = "missioncore.compute-result/v1"
COMPUTE_RESULT_IDENTITY_SCHEMA = "missioncore.compute-result-identity/v1"
OBJECT_DETECTIONS_SCHEMA = "missioncore.object-detections/v1"
SESSION_TIMELINE = "session_time"

MAX_RESULT_JSON_BYTES = 32 * 1024 * 1024
MAX_RESULT_FRAMES = 10_000
MAX_DETECTIONS_PER_FRAME = 10_000
MAX_OVERLAY_SOURCE_BYTES = 512 * 1024 * 1024
MAX_OVERLAY_DECODED_BYTES = 512 * 1024 * 1024
MAX_OVERLAY_RRD_BYTES = 512 * 1024 * 1024
MAX_JOB_SCAN = 512

_SAFE_RESULT_ID = re.compile(r"^result-[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 ObjectDetection:
    class_id: int
    class_name: str
    score: float
    bbox_xyxy: tuple[float, float, float, float]


@dataclass(frozen=True, slots=True)
class DetectionFrame:
    frame_index: int
    session_time_ns: int
    detections: tuple[ObjectDetection, ...]


@dataclass(frozen=True, slots=True)
class RecordedPerceptionResult:
    result_id: str
    result_root: Path
    job: CameraComputeJob
    created_at_utc: str
    frames: tuple[DetectionFrame, ...]


class RecordedPerceptionOverlayError(RuntimeError):
    """A validated compute result could not be projected into Rerun."""


def validate_recorded_perception_result(
    job_root: Path,
    result_root: Path,
) -> RecordedPerceptionResult:
    """Validate one content-addressed worker result against its exact job."""

    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("compute result root is invalid")

    result = _read_json_object(root / "result.json", root)
    detections = _read_json_object(root / "detections.json", root)
    identity_sha256 = result.get("identity_sha256")
    pipeline = result.get("pipeline")
    model = result.get("model")
    parameters = result.get("parameters")
    if (
        result.get("schema_version") != COMPUTE_RESULT_SCHEMA
        or result.get("result_id") != root.name
        or not isinstance(identity_sha256, str)
        or _SHA256.fullmatch(identity_sha256) is None
        or root.name != f"result-{identity_sha256}"
        or result.get("job_id") != job.job_id
        or result.get("input_sha256") != job.input_sha256
        or not isinstance(pipeline, dict)
        or not isinstance(model, dict)
        or not isinstance(parameters, dict)
    ):
        raise SessionIntegrityError("compute result identity is inconsistent")
    identity = {
        "schema_version": COMPUTE_RESULT_IDENTITY_SCHEMA,
        "job_id": job.job_id,
        "input_sha256": job.input_sha256,
        "pipeline": pipeline,
        "model": {
            "id": model.get("id"),
            "version": model.get("version"),
            "sha256": model.get("sha256"),
        },
        "parameters": parameters,
    }
    if hashlib.sha256(_canonical_json(identity)).hexdigest() != identity_sha256:
        raise SessionIntegrityError("compute result generation is inconsistent")

    artifacts = result.get("artifacts")
    detection_path = root / "detections.json"
    detection_stat = _confined_regular_file(detection_path, root)
    if not isinstance(artifacts, list) or len(artifacts) != 1:
        raise SessionIntegrityError("compute result artifact descriptor is invalid")
    artifact = artifacts[0]
    if (
        not isinstance(artifact, dict)
        or artifact.get("kind") != "object-detections"
        or artifact.get("path") != "detections.json"
        or artifact.get("schema_version") != OBJECT_DETECTIONS_SCHEMA
        or artifact.get("byte_length") != detection_stat.st_size
        or artifact.get("sha256") != _sha256_file(detection_path)
    ):
        raise SessionIntegrityError("compute result artifact identity changed")

    if (
        detections.get("schema_version") != OBJECT_DETECTIONS_SCHEMA
        or detections.get("result_id") != root.name
        or detections.get("job_id") != job.job_id
        or detections.get("input_sha256") != job.input_sha256
        or detections.get("timestamp_basis") != "session-time-seconds"
    ):
        raise SessionIntegrityError("object detections are not bound to the result")
    raw_frames = detections.get("frames")
    if not isinstance(raw_frames, list) or not 1 <= len(raw_frames) <= MAX_RESULT_FRAMES:
        raise SessionIntegrityError("object detection frame set is outside bounds")

    frames: list[DetectionFrame] = []
    detection_count = 0
    previous_time_ns = -1
    for expected_index, value in enumerate(raw_frames):
        frame = _validate_detection_frame(value, expected_index, job)
        if frame.session_time_ns <= previous_time_ns:
            raise SessionIntegrityError("object detection timestamps are not increasing")
        previous_time_ns = frame.session_time_ns
        frames.append(frame)
        detection_count += len(frame.detections)

    metrics = result.get("metrics")
    if (
        not isinstance(metrics, dict)
        or metrics.get("frames_processed") != len(frames)
        or metrics.get("detections") != detection_count
    ):
        raise SessionIntegrityError("compute result metrics do not match its artifact")
    created_at_utc = result.get("created_at_utc")
    if not isinstance(created_at_utc, str) or len(created_at_utc) > 64:
        raise SessionIntegrityError("compute result creation time is invalid")
    return RecordedPerceptionResult(
        result_id=root.name,
        result_root=root,
        job=job,
        created_at_utc=created_at_utc,
        frames=tuple(frames),
    )


class RecordedPerceptionOverlayStore:
    """Discover admitted local results and cache complete overlay RRD files."""

    def __init__(
        self,
        *,
        jobs_root: Path,
        results_root: Path,
        cache_root: Path,
        ffmpeg_path: Path,
        ffprobe_path: Path,
    ) -> None:
        self.jobs_root = jobs_root.expanduser().absolute()
        self.results_root = results_root.expanduser().absolute()
        self.cache_root = cache_root.expanduser().absolute()
        self.ffmpeg_path = ffmpeg_path.expanduser().resolve(strict=True)
        self.ffprobe_path = ffprobe_path.expanduser().resolve(strict=True)
        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 perception application id is invalid")
        if _SAFE_RECORDING_ID.fullmatch(recording_id) is None:
            raise ValueError("recorded perception recording id is invalid")
        with self._lock:
            result = self._latest_result(session_id)
            if result is None:
                return None
            cache_root = _private_directory(self.cache_root)
            session_cache = _private_child_directory(cache_root, session_id)
            cache_dir = _private_child_directory(session_cache, result.result_id)
            output = cache_dir / f"{recording_id}.rrd"
            sidecar = output.with_suffix(".rrd.cache.json")
            cached = _read_cached_overlay(output, sidecar, result)
            if cached is not None:
                return cached
            payload = _render_overlay(
                result,
                application_id=application_id,
                recording_id=recording_id,
                ffmpeg_path=self.ffmpeg_path,
                ffprobe_path=self.ffprobe_path,
            )
            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.perception-overlay-cache/v1",
                        "result_id": result.result_id,
                        "recording_id": recording_id,
                        "byte_length": len(payload),
                        "sha256": hashlib.sha256(payload).hexdigest(),
                    },
                )
            finally:
                temporary.unlink(missing_ok=True)
            return payload

    def _latest_result(self, session_id: str) -> RecordedPerceptionResult | None:
        try:
            job_roots = sorted(self.jobs_root.iterdir())
        except FileNotFoundError:
            return None
        if len(job_roots) > MAX_JOB_SCAN:
            raise RecordedPerceptionOverlayError("compute job catalog is outside bounds")
        matches: list[RecordedPerceptionResult] = []
        for job_root in job_roots:
            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
            result_parent = self.results_root / job.job_id
            try:
                result_roots = sorted(
                    path
                    for path in result_parent.iterdir()
                    if path.is_dir()
                    and not path.is_symlink()
                    and _SAFE_RESULT_ID.fullmatch(path.name) is not None
                )
            except FileNotFoundError:
                continue
            for result_root in result_roots:
                matches.append(validate_recorded_perception_result(job_root, result_root))
        if not matches:
            return None
        return max(matches, key=lambda value: (value.created_at_utc, value.result_id))


def _validate_detection_frame(
    value: object,
    expected_index: int,
    job: CameraComputeJob,
) -> DetectionFrame:
    if not isinstance(value, dict) or value.get("frame_index") != expected_index:
        raise SessionIntegrityError("object detection frame index is inconsistent")
    session_seconds = value.get("session_seconds")
    epoch_seconds = value.get("epoch_seconds")
    raw_detections = value.get("detections")
    if (
        not _finite_number(session_seconds)
        or not _finite_number(epoch_seconds)
        or float(epoch_seconds) < 0
        or float(session_seconds) < job.timeline_start_seconds - 0.001
        or float(session_seconds) > job.timeline_end_seconds + 0.001
        or not isinstance(raw_detections, list)
        or len(raw_detections) > MAX_DETECTIONS_PER_FRAME
    ):
        raise SessionIntegrityError("object detection frame is invalid")
    found = tuple(_validate_detection(item) for item in raw_detections)
    return DetectionFrame(
        frame_index=expected_index,
        session_time_ns=round(float(session_seconds) * 1_000_000_000),
        detections=found,
    )


def _validate_detection(value: object) -> ObjectDetection:
    if not isinstance(value, dict):
        raise SessionIntegrityError("object detection is invalid")
    class_id = value.get("class_id")
    class_name = value.get("class_name")
    score = value.get("score")
    bbox = value.get("bbox_xyxy")
    if (
        not isinstance(class_id, int)
        or isinstance(class_id, bool)
        or not 0 <= class_id <= 10_000
        or not isinstance(class_name, str)
        or not class_name
        or len(class_name) > 128
        or not _finite_number(score)
        or not 0 <= float(score) <= 1
        or not isinstance(bbox, list)
        or len(bbox) != 4
        or not all(_finite_number(item) for item in bbox)
    ):
        raise SessionIntegrityError("object detection fields are invalid")
    x1, y1, x2, y2 = (float(item) for item in bbox)
    if min(x1, y1) < 0 or x2 < x1 or y2 < y1 or max(x2, y2) > 100_000:
        raise SessionIntegrityError("object detection bounds are invalid")
    return ObjectDetection(class_id, class_name, float(score), (x1, y1, x2, y2))


def _render_overlay(
    result: RecordedPerceptionResult,
    *,
    application_id: str,
    recording_id: str,
    ffmpeg_path: Path,
    ffprobe_path: Path,
) -> bytes:
    if validate_camera_compute_job(result.job.job_root) != result.job:
        raise RecordedPerceptionOverlayError("camera compute job changed before projection")
    source = _camera_source_bytes(result.job)
    width, height, frame_count = _probe_video(source, ffprobe_path)
    if frame_count != len(result.frames):
        raise RecordedPerceptionOverlayError("camera and detection frame counts differ")
    _validate_image_bounds(result, width, height)
    frame_bytes = width * height * 3
    expected_bytes = frame_bytes * len(result.frames)
    if expected_bytes > MAX_OVERLAY_DECODED_BYTES:
        raise RecordedPerceptionOverlayError("decoded camera overlay is outside bounds")
    decoded = _decode_rgb(source, ffmpeg_path, frame_count)
    if len(decoded) != expected_bytes:
        raise RecordedPerceptionOverlayError("decoded camera frame count changed")

    recording = rr.RecordingStream(
        application_id,
        recording_id=recording_id,
        send_properties=False,
    )
    stream = rr.binary_stream(recording)
    try:
        recording.log(
            "/perception/camera/metadata",
            rr.AnyValues(
                result_id=result.result_id,
                job_id=result.job.job_id,
                source_id=result.job.source_id,
                pipeline="recorded-camera-coco-detection/v1",
                warning="Generic detections are evidence, not a driving decision.",
            ),
            static=True,
        )
        for frame in result.frames:
            offset = frame.frame_index * frame_bytes
            image = np.frombuffer(
                decoded,
                dtype=np.uint8,
                count=frame_bytes,
                offset=offset,
            ).reshape((height, width, 3))
            recording.set_time(
                SESSION_TIMELINE,
                duration=np.timedelta64(frame.session_time_ns, "ns"),
            )
            recording.log(
                "/perception/camera/image",
                rr.Image(image, color_model="RGB"),
            )
            if frame.detections:
                recording.log(
                    "/perception/camera/detections",
                    rr.Boxes2D(
                        array=[detection.bbox_xyxy for detection in frame.detections],
                        array_format=rr.Box2DFormat.XYXY,
                        labels=[
                            f"{detection.class_name} · {detection.score:.0%}"
                            for detection in frame.detections
                        ],
                        colors=[
                            _class_color(detection.class_id)
                            for detection in frame.detections
                        ],
                        show_labels=True,
                    ),
                )
            else:
                recording.log(
                    "/perception/camera/detections",
                    rr.Clear(recursive=False),
                )
        payload = stream.read(flush=True, flush_timeout_sec=120.0)
    except Exception as exc:
        raise RecordedPerceptionOverlayError("failed to serialize perception overlay") from exc
    finally:
        recording.disconnect()
    if payload is None or not payload.startswith(b"RRF2") or len(payload) > MAX_OVERLAY_RRD_BYTES:
        raise RecordedPerceptionOverlayError("serialized perception overlay is invalid")
    return payload


def _camera_source_bytes(job: CameraComputeJob) -> bytes:
    epoch = (
        job.job_root
        / "input"
        / "camera"
        / job.source_id
        / f"epoch-{job.codec_epoch}"
    )
    paths = [epoch / "init.mp4"] + [
        epoch / "segments" / f"{sequence}.m4s"
        for sequence in range(1, job.segment_count + 1)
    ]
    total = sum(path.stat().st_size for path in paths)
    if total > MAX_OVERLAY_SOURCE_BYTES:
        raise RecordedPerceptionOverlayError("camera overlay source is outside bounds")
    return b"".join(path.read_bytes() for path in paths)


def _probe_video(source: bytes, ffprobe_path: Path) -> tuple[int, int, int]:
    completed = _run_media_tool(
        [
            str(ffprobe_path),
            "-v",
            "error",
            "-count_frames",
            "-select_streams",
            "v:0",
            "-show_entries",
            "stream=width,height,nb_read_frames",
            "-of",
            "json",
            "-i",
            "pipe:0",
        ],
        source,
    )
    try:
        document = json.loads(completed.stdout)
        stream = document["streams"][0]
        width = int(stream["width"])
        height = int(stream["height"])
        frame_count = int(stream["nb_read_frames"])
    except (IndexError, KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
        raise RecordedPerceptionOverlayError("camera dimensions are unavailable") from exc
    if (
        width < 1
        or height < 1
        or width * height > 4_194_304
        or not 1 <= frame_count <= MAX_RESULT_FRAMES
    ):
        raise RecordedPerceptionOverlayError("camera dimensions are outside bounds")
    return width, height, frame_count


def _decode_rgb(source: bytes, ffmpeg_path: Path, frame_count: int) -> bytes:
    return _run_media_tool(
        [
            str(ffmpeg_path),
            "-v",
            "error",
            "-i",
            "pipe:0",
            "-map",
            "0:v:0",
            "-frames:v",
            str(frame_count),
            "-fps_mode",
            "passthrough",
            "-pix_fmt",
            "rgb24",
            "-f",
            "rawvideo",
            "pipe:1",
        ],
        source,
    ).stdout


def _run_media_tool(argv: list[str], payload: bytes) -> subprocess.CompletedProcess[bytes]:
    try:
        metadata = Path(argv[0]).lstat()
        if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISREG(metadata.st_mode):
            raise RecordedPerceptionOverlayError("media tool is not a regular file")
        completed = subprocess.run(
            argv,
            input=payload,
            capture_output=True,
            check=False,
            timeout=120,
        )
    except (OSError, subprocess.SubprocessError) as exc:
        raise RecordedPerceptionOverlayError("media tool failed") from exc
    if completed.returncode != 0:
        raise RecordedPerceptionOverlayError("camera media could not be decoded")
    return completed


def _read_cached_overlay(
    output: Path,
    sidecar: Path,
    result: RecordedPerceptionResult,
) -> bytes | None:
    try:
        value = _read_json_object(sidecar, sidecar.parent)
        _confined_regular_file(output, output.parent)
        payload = output.read_bytes()
    except (OSError, SessionIntegrityError):
        return None
    if (
        value.get("schema_version") != "missioncore.perception-overlay-cache/v1"
        or value.get("result_id") != result.result_id
        or value.get("recording_id") != output.stem
        or value.get("byte_length") != len(payload)
        or value.get("sha256") != hashlib.sha256(payload).hexdigest()
        or not payload.startswith(b"RRF2")
        or len(payload) > MAX_OVERLAY_RRD_BYTES
    ):
        return None
    return payload


def _validate_image_bounds(
    result: RecordedPerceptionResult,
    width: int,
    height: int,
) -> None:
    for frame in result.frames:
        for detection in frame.detections:
            x1, y1, x2, y2 = detection.bbox_xyxy
            if x1 > width or x2 > width or y1 > height or y2 > height:
                raise RecordedPerceptionOverlayError(
                    "object detection escapes the decoded camera image"
                )


def _private_directory(path: Path) -> Path:
    path.mkdir(mode=0o700, parents=True, exist_ok=True)
    try:
        metadata = path.lstat()
    except OSError as exc:
        raise RecordedPerceptionOverlayError("perception cache is unavailable") from exc
    if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISDIR(metadata.st_mode):
        raise RecordedPerceptionOverlayError("perception cache is not a real directory")
    resolved = path.resolve(strict=True)
    os.chmod(resolved, 0o700)
    return resolved


def _private_child_directory(parent: Path, name: str) -> Path:
    child = parent / name
    child.mkdir(mode=0o700, exist_ok=True)
    try:
        metadata = child.lstat()
        resolved = child.resolve(strict=True)
    except OSError as exc:
        raise RecordedPerceptionOverlayError("perception cache child is unavailable") from exc
    if (
        stat.S_ISLNK(metadata.st_mode)
        or not stat.S_ISDIR(metadata.st_mode)
        or resolved.parent != parent
    ):
        raise RecordedPerceptionOverlayError("perception cache child escapes its root")
    os.chmod(resolved, 0o700)
    return resolved


def _class_color(class_id: int) -> list[int]:
    palette = (
        (185, 255, 74, 255),
        (67, 191, 255, 255),
        (255, 193, 92, 255),
        (222, 110, 255, 255),
        (255, 103, 117, 255),
    )
    return list(palette[class_id % len(palette)])


def _read_json_object(path: Path, root: Path) -> dict[str, Any]:
    metadata = _confined_regular_file(path, root)
    if not 0 < metadata.st_size <= MAX_RESULT_JSON_BYTES:
        raise SessionIntegrityError("compute result JSON is outside bounds")
    try:
        value = json.loads(path.read_text(encoding="utf-8-sig"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise SessionIntegrityError("compute result JSON is unavailable") from exc
    if not isinstance(value, dict):
        raise SessionIntegrityError("compute result 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("compute result 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("compute result artifact is not a confined regular file")
    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("compute result identity cannot be encoded") 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))
    )
