"""Publish zero-copy replay caches for immutable LAB session aliases."""

from __future__ import annotations

import hashlib
import json
import os
import shutil
import stat
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any

from k1link.artifacts import write_json_atomic

from .models import SessionIntegrityError
from .recording import (
    COMPATIBLE_CACHE_SCHEMAS,
    RECORDING_CACHE_FILENAME,
    RECORDING_CACHE_SIDECAR_FILENAME,
)


def publish_lab_replay_cache(
    data_dir: Path,
    *,
    source_session_id: str,
    lab_session_id: str,
    timeline_start_ns: int | None = None,
    timeline_end_ns: int | None = None,
) -> None:
    """Publish a validated source RRD and clone its path-free sidecars.

    Full-session LAB instances hard-link the immutable source bytes. Bounded
    experiments materialize only their declared timeline window, so a short AI
    run cannot silently hold its final frame over the rest of a long source
    recording.
    """

    bounded = timeline_start_ns is not None or timeline_end_ns is not None
    if bounded and (
        not isinstance(timeline_start_ns, int)
        or isinstance(timeline_start_ns, bool)
        or not isinstance(timeline_end_ns, int)
        or isinstance(timeline_end_ns, bool)
        or timeline_start_ns < 0
        or timeline_end_ns <= timeline_start_ns
    ):
        raise ValueError("LAB replay window is invalid")
    window_start_ns = -1
    window_end_ns = -1
    if bounded:
        assert isinstance(timeline_start_ns, int)
        assert isinstance(timeline_end_ns, int)
        window_start_ns = timeline_start_ns
        window_end_ns = timeline_end_ns

    root = data_dir.expanduser().resolve(strict=True)
    recordings = (root / "recordings").resolve(strict=True)
    source_root = recordings / source_session_id
    source_rrd = source_root / RECORDING_CACHE_FILENAME
    source_sidecar = source_root / RECORDING_CACHE_SIDECAR_FILENAME
    _require_regular(source_rrd, source_root)
    _require_regular(source_sidecar, source_root)
    document = _read_object(source_sidecar)
    if (
        document.get("schema_version") not in COMPATIBLE_CACHE_SCHEMAS
        or document.get("session_id") != source_session_id
    ):
        raise SessionIntegrityError("source recording cache is incompatible")

    target_root = recordings / lab_session_id
    target_root.mkdir(mode=0o700, parents=False, exist_ok=True)
    if target_root.is_symlink() or target_root.resolve() != target_root:
        raise SessionIntegrityError("LAB recording cache root is unsafe")
    target_rrd = target_root / RECORDING_CACHE_FILENAME
    if target_rrd.exists():
        if bounded:
            _require_matching_bounded_cache(
                target_rrd,
                target_root / RECORDING_CACHE_SIDECAR_FILENAME,
                lab_session_id=lab_session_id,
                timeline_start_ns=window_start_ns,
                timeline_end_ns=window_end_ns,
            )
        else:
            source_stat = _require_regular(source_rrd, source_root)
            target_stat = _require_regular(target_rrd, target_root)
            if (source_stat.st_dev, source_stat.st_ino) != (
                target_stat.st_dev,
                target_stat.st_ino,
            ):
                raise SessionIntegrityError(
                    "LAB recording cache already contains different bytes"
                )
    elif bounded:
        _publish_bounded_rrd(
            source_rrd,
            target_rrd,
            recording_id=lab_session_id,
            timeline_start_ns=window_start_ns,
            timeline_end_ns=window_end_ns,
        )
    else:
        temporary = target_root / f".{RECORDING_CACHE_FILENAME}.{os.getpid()}.tmp"
        try:
            os.link(source_rrd, temporary, follow_symlinks=False)
            os.replace(temporary, target_rrd)
        finally:
            temporary.unlink(missing_ok=True)
    target_stat = _require_regular(target_rrd, target_root)
    cloned = {
        **document,
        "session_id": lab_session_id,
        "recording_byte_length": target_stat.st_size,
        "recording_mtime_ns": target_stat.st_mtime_ns,
        "recording_sha256": _sha256(target_rrd),
        "timeline_start_ns": (
            window_start_ns if bounded else document["timeline_start_ns"]
        ),
        "timeline_end_ns": window_end_ns if bounded else document["timeline_end_ns"],
    }
    write_json_atomic(target_root / RECORDING_CACHE_SIDECAR_FILENAME, cloned)
    os.chmod(target_root / RECORDING_CACHE_SIDECAR_FILENAME, 0o600)
    if not bounded:
        _clone_recorded_media_sidecars(
            root / "recorded-media-preparations",
            source_session_id=source_session_id,
            lab_session_id=lab_session_id,
        )


def _publish_bounded_rrd(
    source_rrd: Path,
    target_rrd: Path,
    *,
    recording_id: str,
    timeline_start_ns: int,
    timeline_end_ns: int,
) -> None:
    rerun = _rerun_cli()
    end_exclusive_ns = timeline_end_ns + 1
    split_recording_prefix = f"{recording_id}-window-"
    split_recording_id = f"{split_recording_prefix}1"
    with tempfile.TemporaryDirectory(
        prefix=".lab-window-",
        dir=target_rrd.parent,
    ) as temporary_name:
        temporary_root = Path(temporary_name)
        split_root = temporary_root / "split"
        split_root.mkdir()
        _run_rerun(
            rerun,
            "rrd",
            "split",
            "--output-dir",
            str(split_root),
            "--timeline",
            "session_time",
            "--time",
            str(timeline_start_ns),
            "--time",
            str(end_exclusive_ns),
            "--recording-id",
            split_recording_prefix,
            str(source_rrd),
        )
        split_name = (
            f"{source_rrd.stem}_{_rrd_duration(timeline_start_ns)}"
            f"__{_rrd_duration(end_exclusive_ns)}.rrd"
        )
        split_path = split_root / split_name
        _require_regular(split_path, split_root)
        marker = temporary_root / "window.rrd"
        _write_window_markers(
            marker,
            recording_id=split_recording_id,
            timeline_start_ns=timeline_start_ns,
            timeline_end_ns=timeline_end_ns,
        )
        merged = temporary_root / "bounded.rrd"
        _run_rerun(
            rerun,
            "rrd",
            "merge",
            "--output",
            str(merged),
            str(split_path),
            str(marker),
        )
        _run_rerun(rerun, "rrd", "verify", str(merged))
        _require_regular(merged, temporary_root)
        temporary_target = target_rrd.parent / f".{target_rrd.name}.{os.getpid()}.tmp"
        try:
            os.replace(merged, temporary_target)
            os.chmod(temporary_target, 0o600)
            os.replace(temporary_target, target_rrd)
        finally:
            temporary_target.unlink(missing_ok=True)


def _write_window_markers(
    path: Path,
    *,
    recording_id: str,
    timeline_start_ns: int,
    timeline_end_ns: int,
) -> None:
    import numpy as np
    import rerun as rr

    recording = rr.RecordingStream(
        "nodedc_mission_core_recorded",
        recording_id=recording_id,
    )
    stream = rr.binary_stream(recording)
    try:
        for timestamp, label in (
            (timeline_start_ns, "LAB window start"),
            (timeline_end_ns, "LAB window end"),
        ):
            recording.set_time(
                "session_time",
                duration=np.timedelta64(timestamp, "ns"),
            )
            recording.log(
                "/__mission_core/lab_window",
                rr.TextDocument(label),
            )
        payload = stream.read(flush=True, flush_timeout_sec=5.0)
    finally:
        recording.disconnect()
    if payload is None or not payload.startswith(b"RRF2"):
        raise SessionIntegrityError("LAB replay window markers are invalid")
    path.write_bytes(payload)
    os.chmod(path, 0o600)


def _rerun_cli() -> Path:
    adjacent = Path(sys.executable).with_name("rerun")
    candidate = adjacent if adjacent.is_file() else None
    if candidate is None:
        resolved = shutil.which("rerun")
        candidate = Path(resolved) if resolved is not None else None
    if candidate is None or not candidate.is_file() or not os.access(candidate, os.X_OK):
        raise SessionIntegrityError("Rerun CLI is unavailable for bounded LAB replay")
    return candidate.resolve()


def _run_rerun(executable: Path, *arguments: str) -> None:
    try:
        completed = subprocess.run(
            [str(executable), *arguments],
            check=False,
            capture_output=True,
            text=True,
            timeout=120,
        )
    except (OSError, subprocess.TimeoutExpired) as exc:
        raise SessionIntegrityError("bounded LAB replay command failed") from exc
    if completed.returncode != 0:
        detail = completed.stderr.strip()[-1000:]
        raise SessionIntegrityError(f"bounded LAB replay command failed: {detail}")


def _rrd_duration(value_ns: int) -> str:
    for suffix, unit in (
        ("s", 1_000_000_000),
        ("ms", 1_000_000),
        ("us", 1_000),
        ("ns", 1),
    ):
        if value_ns < unit:
            continue
        whole, remainder = divmod(value_ns, unit)
        if remainder == 0:
            return f"{whole}{suffix}"
        digits = len(str(unit)) - 1
        fraction = f"{remainder:0{digits}d}".rstrip("0")
        return f"{whole}_{fraction}{suffix}"
    return "0ns"


def _require_matching_bounded_cache(
    target_rrd: Path,
    target_sidecar: Path,
    *,
    lab_session_id: str,
    timeline_start_ns: int,
    timeline_end_ns: int,
) -> None:
    target_root = target_rrd.parent
    metadata = _require_regular(target_rrd, target_root)
    _require_regular(target_sidecar, target_root)
    document = _read_object(target_sidecar)
    if (
        document.get("session_id") != lab_session_id
        or document.get("recording_byte_length") != metadata.st_size
        or document.get("recording_mtime_ns") != metadata.st_mtime_ns
        or document.get("recording_sha256") != _sha256(target_rrd)
        or document.get("timeline_start_ns") != timeline_start_ns
        or document.get("timeline_end_ns") != timeline_end_ns
    ):
        raise SessionIntegrityError("LAB recording cache already contains different bytes")


def _clone_recorded_media_sidecars(
    root: Path,
    *,
    source_session_id: str,
    lab_session_id: str,
) -> None:
    if not root.is_dir():
        return
    for source in root.iterdir():
        if source.is_symlink() or not source.is_file():
            continue
        try:
            document = _read_object(source)
        except (OSError, SessionIntegrityError):
            continue
        if document.get("session_id") != source_session_id:
            continue
        artifact_id = document.get("artifact_id")
        if not isinstance(artifact_id, str):
            raise SessionIntegrityError("recorded media sidecar has no artifact identity")
        body = {key: value for key, value in document.items() if key != "checksum_sha256"}
        body["session_id"] = lab_session_id
        checksum = hashlib.sha256(_canonical_json(body)).hexdigest()
        target = root / _media_sidecar_name(lab_session_id, artifact_id)
        write_json_atomic(target, {**body, "checksum_sha256": checksum})
        os.chmod(target, 0o600)


def _media_sidecar_name(session_id: str, artifact_id: str) -> str:
    key = f"{session_id}\0{artifact_id}".encode()
    return f"{hashlib.sha256(key).hexdigest()}.json"


def _canonical_json(value: object) -> bytes:
    return json.dumps(
        value,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
        allow_nan=False,
    ).encode("utf-8")


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


def _read_object(path: Path) -> dict[str, Any]:
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise SessionIntegrityError("cache sidecar is not a JSON object")
    return value


def _require_regular(path: Path, parent: Path) -> os.stat_result:
    metadata = path.lstat()
    if (
        stat.S_ISLNK(metadata.st_mode)
        or not stat.S_ISREG(metadata.st_mode)
        or path.parent.resolve(strict=True) != parent.resolve(strict=True)
    ):
        raise SessionIntegrityError("cache artifact is not a confined regular file")
    return metadata
