#!/usr/bin/env python3
"""Source-paced RAVNOVES00 qualification of the live shadow transport."""

from __future__ import annotations

import argparse
import hashlib
import heapq
import itertools
import json
import math
import os
import resource
import statistics
import sys
import threading
import time
from collections.abc import Iterator
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Literal

import uvicorn
from fastapi import FastAPI

from k1link.compute.live_perception import (
    LatestWinsQueue,
    LiveIngressModality,
    LivePerceptionIngress,
    LivePerceptionResultFrame,
    decode_live_perception_result,
)
from k1link.device_plugins.xgrids_k1.live_perception_shadow import (
    build_live_perception_shadow_router,
    ensure_live_shadow_token,
)
from k1link.device_plugins.xgrids_k1.viewer.replay import iter_replay_messages

REPORT_SCHEMA = "missioncore.e12-shadow-transport-source-report/v1"
PLUGIN_ID = "nodedc.device.xgrids-lixelkity-k1"
_EXPERIMENT_ID = "e12"


@dataclass(frozen=True, slots=True)
class ReplayEvent:
    modality: LiveIngressModality
    source_id: str
    source_sequence: int
    epoch_ns: int
    monotonic_ns: int
    payload: bytes


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


def _percentiles(values: list[float]) -> dict[str, float | int]:
    if not values:
        return {
            "count": 0,
            "mean": 0.0,
            "p50": 0.0,
            "p95": 0.0,
            "p99": 0.0,
            "max": 0.0,
        }
    ordered = sorted(values)

    def percentile(fraction: float) -> float:
        index = min(len(ordered) - 1, max(0, round((len(ordered) - 1) * fraction)))
        return round(ordered[index], 6)

    return {
        "count": len(values),
        "mean": round(statistics.fmean(values), 6),
        "p50": percentile(0.50),
        "p95": percentile(0.95),
        "p99": percentile(0.99),
        "max": round(max(values), 6),
    }


def _camera_events(epoch_root: Path, duration_seconds: float) -> Iterator[ReplayEvent]:
    index_path = epoch_root / "index.jsonl"
    if not index_path.is_file():
        raise RuntimeError("camera index is missing")
    with index_path.open("r", encoding="utf-8") as stream:
        first_line = stream.readline()
        if not first_line:
            raise RuntimeError("camera replay interval is empty")
        first_record = json.loads(first_line)
        start_epoch_ns = int(first_record["host_epoch_ns"])
        end_epoch_ns = start_epoch_ns + int(duration_seconds * 1_000_000_000)
        init_payload = (epoch_root / "init.mp4").read_bytes()
        yield ReplayEvent(
            modality="camera-init",
            source_id="sensor.camera.right",
            source_sequence=0,
            epoch_ns=start_epoch_ns - 1,
            monotonic_ns=int(first_record["host_monotonic_ns"]) - 1,
            payload=init_payload,
        )
        for line in itertools.chain((first_line,), stream):
            record = json.loads(line)
            epoch_ns = int(record["host_epoch_ns"])
            if epoch_ns > end_epoch_ns:
                break
            payload_path = epoch_root / str(record["path"])
            payload = payload_path.read_bytes()
            if len(payload) != int(record["length"]):
                raise RuntimeError("camera replay segment length differs from its index")
            if hashlib.sha256(payload).hexdigest() != record["sha256"]:
                raise RuntimeError("camera replay segment digest differs from its index")
            yield ReplayEvent(
                modality="camera-frame",
                source_id="sensor.camera.right",
                source_sequence=int(record["sequence"]),
                epoch_ns=epoch_ns,
                monotonic_ns=int(record["host_monotonic_ns"]),
                payload=payload,
            )


def _mqtt_events(
    capture_path: Path,
    *,
    start_epoch_ns: int,
    duration_seconds: float,
) -> Iterator[ReplayEvent]:
    end_epoch_ns = start_epoch_ns + int(duration_seconds * 1_000_000_000)
    for message in iter_replay_messages(capture_path):
        if message.received_at_epoch_ns < start_epoch_ns:
            continue
        if message.received_at_epoch_ns > end_epoch_ns:
            break
        if message.topic.endswith("/lio_pcl") or message.topic == "RealtimePointcloud":
            modality: Literal["lidar", "pose"] = "lidar"
        elif message.topic.endswith("/lio_pose") or message.topic == "RealtimePath":
            modality = "pose"
        else:
            continue
        yield ReplayEvent(
            modality=modality,
            source_id=message.topic,
            source_sequence=message.sequence,
            epoch_ns=message.received_at_epoch_ns,
            monotonic_ns=message.received_monotonic_ns or 0,
            payload=message.payload,
        )


def _wait_for(predicate: object, *, timeout_seconds: float, label: str) -> None:
    if not callable(predicate):
        raise TypeError("wait predicate must be callable")
    deadline = time.monotonic() + timeout_seconds
    while time.monotonic() < deadline:
        if predicate():
            return
        time.sleep(0.05)
    raise RuntimeError(f"timed out waiting for {label}")


def run(args: argparse.Namespace) -> dict[str, object]:
    experiment_id = str(getattr(args, "experiment_id", _EXPERIMENT_ID))
    speed = float(getattr(args, "speed", 1.0))
    if (
        experiment_id not in {"e12", "e21", "e23"}
        or not math.isfinite(speed)
        or not 0.1 <= speed <= 10
    ):
        raise RuntimeError("replay experiment identity or speed is invalid")
    repository_root = Path(args.repository_root).expanduser().resolve()
    session_root = Path(args.session_root).expanduser().resolve()
    epoch_root = session_root / "media" / "sensor.camera.right" / "epoch-1"
    capture_path = session_root / "captures" / "mqtt_live" / "mqtt.raw.k1mqtt"
    camera = iter(_camera_events(epoch_root, args.duration_seconds))
    first_camera = next(camera, None)
    first_frame = next(camera, None)
    if first_camera is None or first_frame is None or first_frame.modality != "camera-frame":
        raise RuntimeError("camera replay interval is empty")
    camera = itertools.chain((first_camera, first_frame), camera)
    mqtt = iter(
        _mqtt_events(
            capture_path,
            start_epoch_ns=first_frame.epoch_ns,
            duration_seconds=args.duration_seconds,
        )
    )
    first_mqtt = next(mqtt, None)
    if first_mqtt is None:
        raise RuntimeError("selected replay interval contains no LiDAR or pose events")
    mqtt = itertools.chain((first_mqtt,), mqtt)
    events = heapq.merge(
        camera,
        mqtt,
        key=lambda event: (event.epoch_ns, event.modality),
    )

    ingress = LivePerceptionIngress()
    result_queue = LatestWinsQueue[LivePerceptionResultFrame](capacity=2)
    result_state = {"received": 0, "payload_bytes": 0}
    result_lock = threading.Lock()

    def receive_result(payload: bytes) -> bool:
        frame = decode_live_perception_result(payload)
        result_queue.publish(frame)
        with result_lock:
            result_state["received"] += 1
            result_state["payload_bytes"] += len(payload)
        return True

    _, token = ensure_live_shadow_token(repository_root)
    app = FastAPI()
    app.include_router(
        build_live_perception_shadow_router(
            ingress,
            PLUGIN_ID,
            bearer_token=token,
            result_receiver=receive_result,
        )
    )
    server = uvicorn.Server(
        uvicorn.Config(
            app,
            host="127.0.0.1",
            port=args.port,
            log_level="warning",
            access_log=False,
        )
    )
    server_thread = threading.Thread(target=server.run, name="e12-shadow-source", daemon=True)
    server_thread.start()
    _wait_for(lambda: server.started, timeout_seconds=10.0, label="E12 source server")

    session_id = f"{experiment_id}-ravnoves00-{time.time_ns()}"
    ingress.begin_session(session_id)
    _wait_for(
        lambda: ingress.snapshot()["consumer_connected"],
        timeout_seconds=args.consumer_timeout_seconds,
        label="exclusive shadow worker",
    )
    started_monotonic = time.monotonic()
    source_start_ns: int | None = None
    source_last_ns: int | None = None
    published: dict[str, int] = {}
    selected: dict[str, int] = {}
    release_lag_ms: list[float] = []
    publish_call_ms: list[float] = []
    process_cpu_started = time.process_time()
    try:
        for event in events:
            if source_start_ns is None:
                source_start_ns = event.epoch_ns
            source_last_ns = event.epoch_ns
            selected[event.modality] = selected.get(event.modality, 0) + 1
            target = started_monotonic + (
                (event.epoch_ns - source_start_ns) / 1_000_000_000 / speed
            )
            remaining = target - time.monotonic()
            if remaining > 0:
                time.sleep(remaining)
            release_lag_ms.append(max(0.0, (time.monotonic() - target) * 1000))
            publish_started = time.perf_counter()
            accepted = ingress.publish(
                modality=event.modality,
                source_id=event.source_id,
                source_sequence=event.source_sequence,
                captured_at_epoch_ns=event.epoch_ns,
                received_monotonic_ns=event.monotonic_ns,
                payload=event.payload,
            )
            publish_call_ms.append((time.perf_counter() - publish_started) * 1000)
            if accepted:
                published[event.modality] = published.get(event.modality, 0) + 1
        ingress.end_session(session_id)
        _wait_for(
            lambda: not ingress.snapshot()["consumer_connected"],
            timeout_seconds=float(getattr(args, "completion_timeout_seconds", 30.0)),
            label="shadow worker report completion",
        )
    finally:
        ingress.end_session(session_id)
        ingress.close()
        server.should_exit = True
        server_thread.join(timeout=10.0)

    completed_monotonic = time.monotonic()
    source_snapshot = ingress.snapshot()
    if source_start_ns is None or source_last_ns is None:
        raise RuntimeError("selected replay interval is empty")
    source_span_seconds = max(0.0, (source_last_ns - source_start_ns) / 1_000_000_000)
    wall_seconds = completed_monotonic - started_monotonic
    report_schema = {
        "e12": REPORT_SCHEMA,
        "e21": "missioncore.e21-replay-source-report/v1",
        "e23": "missioncore.e23-replay-source-report/v1",
    }[experiment_id]
    peak_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    peak_rss_mib = peak_rss / (1024 * 1024 if sys.platform == "darwin" else 1024)
    report: dict[str, object] = {
        "schema_version": report_schema,
        "state": "completed",
        "session_id": session_id,
        "source": {
            "kind": "recorded-source-paced-near-live",
            "session": session_root.name,
            "camera_source": "sensor.camera.right",
            "camera_epoch": 1,
            "duration_seconds": args.duration_seconds,
            "selected_span_seconds": source_span_seconds,
            "speed": speed,
            "pacing_clock": "recorded-host-arrival",
            "look_ahead": False,
            "payload_loading": "streaming-one-event-per-modality",
            "mqtt_raw_sha256": _sha256(capture_path),
            "camera_index_sha256": _sha256(epoch_root / "index.jsonl"),
        },
        "events_selected": selected,
        "events_admitted": published,
        "wall_seconds": wall_seconds,
        "pacing": {
            "source_span_seconds": source_span_seconds,
            "ideal_wall_seconds": source_span_seconds / speed,
            "wall_over_ideal": wall_seconds / max(source_span_seconds / speed, 1e-9),
            "release_lag_ms": _percentiles(release_lag_ms),
            "publish_call_ms": _percentiles(publish_call_ms),
        },
        "process": {
            "cpu_seconds": time.process_time() - process_cpu_started,
            "peak_rss_mib": peak_rss_mib,
        },
        "diagnostic_results": {
            **result_state,
            "queue": asdict(result_queue.snapshot()),
        },
        "ingress": source_snapshot,
        "authority": {
            "mode": "shadow-diagnostic-only",
            "commands_enabled": False,
            "navigation_or_safety_accepted": False,
        },
    }
    output_root = repository_root / ".runtime" / "compute-experiments" / experiment_id
    output_root.mkdir(parents=True, exist_ok=True)
    stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    temporary = output_root / f".{stamp}-{experiment_id}-source.json.tmp"
    destination = output_root / f"{stamp}-{experiment_id}-source.json"
    encoded = json.dumps(report, indent=2, sort_keys=True).encode("utf-8") + b"\n"
    with temporary.open("xb") as stream:
        stream.write(encoded)
        stream.flush()
        os.fsync(stream.fileno())
    temporary.replace(destination)
    report["report_path"] = str(destination)
    return report


def _arguments() -> argparse.Namespace:
    root = Path(__file__).resolve().parents[2]
    parser = argparse.ArgumentParser()
    parser.add_argument("--repository-root", default=str(root))
    parser.add_argument(
        "--session-root",
        default=str(
            root
            / ".runtime"
            / "mission-core"
            / "evidence"
            / "sessions"
            / "20260720T065719Z_viewer_live"
        ),
    )
    parser.add_argument("--port", type=int, default=8012)
    parser.add_argument("--duration-seconds", type=float, default=15.0)
    parser.add_argument("--speed", type=float, default=1.0)
    parser.add_argument("--experiment-id", choices=("e12", "e21", "e23"), default="e12")
    parser.add_argument("--consumer-timeout-seconds", type=float, default=30.0)
    parser.add_argument("--completion-timeout-seconds", type=float, default=30.0)
    return parser.parse_args()


if __name__ == "__main__":
    print(json.dumps(run(_arguments()), indent=2, sort_keys=True))
