#!/usr/bin/env python3
"""Review and seal the LAB E2 recorded-perception evaluation pack."""

from __future__ import annotations

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

from PIL import Image, ImageDraw, ImageFont

import k1link.compute.evaluation_pack as evaluation_pack_module
from k1link.artifacts import utc_now_iso, write_json_atomic
from k1link.compute import (
    EvaluationFrameRequest,
    prepare_recorded_evaluation_pack,
    validate_camera_compute_job,
    validate_recorded_qualification_slice,
)

CANDIDATE_SCHEMA = "missioncore.e2-evaluation-candidates/v1"
SELECTION_SCHEMA = "missioncore.e2-evaluation-selection/v1"


def _arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    commands = parser.add_subparsers(dest="command", required=True)

    candidates = commands.add_parser("candidates")
    candidates.add_argument("--job-root", type=Path, required=True)
    candidates.add_argument("--qualification-root", type=Path, required=True)
    candidates.add_argument("--output-root", type=Path, required=True)

    seal = commands.add_parser("seal")
    seal.add_argument("--job-root", type=Path, required=True)
    seal.add_argument("--qualification-root", type=Path, required=True)
    seal.add_argument("--valid-fov-root", type=Path, required=True)
    seal.add_argument("--selection", type=Path, required=True)
    seal.add_argument("--output-root", type=Path, required=True)
    return parser.parse_args()


def _run(command: list[str]) -> subprocess.CompletedProcess[str]:
    return subprocess.run(command, check=True, text=True, capture_output=True)


def _ffmpeg_version() -> str:
    completed = _run(["ffmpeg", "-version"])
    return completed.stdout.splitlines()[0].strip()


def _reconstruct_stream(job_root: Path, destination: Path) -> None:
    job = validate_camera_compute_job(job_root)
    epoch = (
        job.job_root
        / "input"
        / "camera"
        / job.source_id
        / f"epoch-{job.codec_epoch}"
    )
    with destination.open("xb") as output:
        with (epoch / "init.mp4").open("rb") as source:
            shutil.copyfileobj(source, output, 1024 * 1024)
        for sequence in range(1, job.segment_count + 1):
            with (epoch / "segments" / f"{sequence}.m4s").open("rb") as source:
                shutil.copyfileobj(source, output, 1024 * 1024)
        output.flush()
        os.fsync(output.fileno())


def _timeline(stream_path: Path, start_seconds: float, expected_count: int) -> list[float]:
    completed = _run(
        [
            "ffprobe",
            "-v",
            "error",
            "-select_streams",
            "v:0",
            "-show_entries",
            "frame=best_effort_timestamp_time",
            "-of",
            "json",
            str(stream_path),
        ]
    )
    document = json.loads(completed.stdout)
    frames = document.get("frames")
    if not isinstance(frames, list) or len(frames) != expected_count:
        raise RuntimeError("decoded camera timestamp count differs from the compute job")
    epoch_values = [float(frame["best_effort_timestamp_time"]) for frame in frames]
    first = epoch_values[0]
    values = [start_seconds + value - first for value in epoch_values]
    if any(right <= left for left, right in zip(values, values[1:], strict=False)):
        raise RuntimeError("decoded camera timestamps are not strictly monotonic")
    return values


def _extract_frames(stream_path: Path, indices: list[int], output_root: Path) -> None:
    if not indices or indices != sorted(set(indices)):
        raise RuntimeError("frame extraction indices must be sorted and unique")
    output_root.mkdir(mode=0o700, parents=True, exist_ok=False)
    expression = "+".join(f"eq(n\\,{index})" for index in indices)
    pattern = output_root / "selected-%06d.png"
    _run(
        [
            "ffmpeg",
            "-hide_banner",
            "-loglevel",
            "error",
            "-i",
            str(stream_path),
            "-map",
            "0:v:0",
            "-vf",
            f"select={expression}",
            "-fps_mode",
            "passthrough",
            str(pattern),
        ]
    )
    extracted = sorted(output_root.glob("selected-*.png"))
    if len(extracted) != len(indices):
        raise RuntimeError("ffmpeg did not extract the requested frame set")
    for source, frame_index in zip(extracted, indices, strict=True):
        destination = output_root / f"frame-{frame_index:06d}.png"
        source.replace(destination)
        os.chmod(destination, 0o600)


def _sha256(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 _contact_sheets(
    frames_root: Path,
    indices: list[int],
    timestamps: list[float],
    output_root: Path,
) -> list[dict[str, Any]]:
    output_root.mkdir(mode=0o700, parents=True, exist_ok=False)
    columns = 8
    rows = 4
    thumb_size = (200, 150)
    label_height = 25
    font = ImageFont.load_default(size=15)
    records: list[dict[str, Any]] = []
    page_size = columns * rows
    for page, offset in enumerate(range(0, len(indices), page_size), start=1):
        page_indices = indices[offset : offset + page_size]
        canvas = Image.new(
            "RGB",
            (columns * thumb_size[0], rows * (thumb_size[1] + label_height)),
            (18, 18, 18),
        )
        draw = ImageDraw.Draw(canvas)
        for position, frame_index in enumerate(page_indices):
            row, column = divmod(position, columns)
            x = column * thumb_size[0]
            y = row * (thumb_size[1] + label_height)
            with Image.open(frames_root / f"frame-{frame_index:06d}.png") as opened:
                image = opened.convert("RGB")
            image.thumbnail(thumb_size, Image.Resampling.LANCZOS)
            canvas.paste(image, (x, y))
            draw.text(
                (x + 4, y + thumb_size[1] + 3),
                f"f={frame_index}  t={timestamps[frame_index]:.3f}s",
                fill=(240, 240, 240),
                font=font,
            )
        path = output_root / f"contact-sheet-{page:02d}.png"
        canvas.save(path, format="PNG", optimize=False)
        os.chmod(path, 0o600)
        records.append(
            {
                "path": path.name,
                "first_frame_index": page_indices[0],
                "last_frame_index": page_indices[-1],
                "frame_count": len(page_indices),
                "byte_length": path.stat().st_size,
                "sha256": _sha256(path),
            }
        )
    return records


def _candidates(args: argparse.Namespace) -> int:
    job = validate_camera_compute_job(args.job_root)
    qualification = validate_recorded_qualification_slice(
        args.qualification_root,
        job_root=job.job_root,
    )
    output = args.output_root.expanduser()
    output.mkdir(mode=0o700, parents=True, exist_ok=False)
    published = False
    try:
        with tempfile.TemporaryDirectory(prefix="missioncore-e2-") as temporary:
            stream_path = Path(temporary) / "camera.mp4"
            _reconstruct_stream(job.job_root, stream_path)
            timestamps = _timeline(stream_path, job.timeline_start_seconds, job.segment_count)
            indices = [frame.frame_index for frame in qualification.frames]
            frames_root = output / "frames"
            _extract_frames(stream_path, indices, frames_root)
        sheets = _contact_sheets(frames_root, indices, timestamps, output / "contact-sheets")
        manifest = {
            "schema_version": CANDIDATE_SCHEMA,
            "created_at_utc": utc_now_iso(),
            "state": "review-candidates",
            "job_id": job.job_id,
            "input_sha256": job.input_sha256,
            "qualification_generation_id": qualification.generation_id,
            "ffmpeg": _ffmpeg_version(),
            "frames": [
                {
                    "qualification_order": order,
                    "frame_index": frame.frame_index,
                    "sequence": frame.sequence,
                    "segment_sha256": frame.segment_sha256,
                    "session_seconds": timestamps[frame.frame_index],
                    "path": f"frames/frame-{frame.frame_index:06d}.png",
                    "byte_length": (
                        frames_root / f"frame-{frame.frame_index:06d}.png"
                    ).stat().st_size,
                    "sha256": _sha256(frames_root / f"frame-{frame.frame_index:06d}.png"),
                }
                for order, frame in enumerate(qualification.frames, start=1)
            ],
            "contact_sheets": sheets,
            "selection_guidance": {
                "anchor_target": 48,
                "temporal_clip_target": {"clips": 4, "frames_per_clip": 4},
                "total_target": 64,
                "required_coverage": [
                    "person",
                    "car_or_heavy_vehicle",
                    "building_structure",
                    "paved_or_unpaved_ground",
                    "grass_and_woody_vegetation",
                    "lens_boundary_hard_negative",
                    "motion_or_occlusion",
                ],
            },
        }
        write_json_atomic(output / "manifest.json", manifest)
        os.chmod(output / "manifest.json", 0o600)
        published = True
    finally:
        if not published and output.exists():
            shutil.rmtree(output)
    print(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True))
    return 0


def _load_selection(path: Path) -> tuple[EvaluationFrameRequest, ...]:
    document = json.loads(path.expanduser().read_text(encoding="utf-8"))
    rows = document.get("frames") if isinstance(document, dict) else None
    if document.get("schema_version") != SELECTION_SCHEMA or not isinstance(rows, list):
        raise RuntimeError("E2 selection document is incompatible")
    selection: list[EvaluationFrameRequest] = []
    for row in rows:
        if not isinstance(row, dict):
            raise RuntimeError("E2 selection frame is invalid")
        selection.append(
            EvaluationFrameRequest(
                frame_index=int(row["frame_index"]),
                role=str(row["role"]),
                group_id=str(row["group_id"]),
            )
        )
    return tuple(selection)


def _write_timeline(path: Path, timestamps: list[float]) -> None:
    with path.open("x", encoding="utf-8", newline="\n") as stream:
        for frame_index, session_seconds in enumerate(timestamps):
            stream.write(
                json.dumps(
                    {"frame_index": frame_index, "session_seconds": session_seconds},
                    separators=(",", ":"),
                )
                + "\n"
            )
        stream.flush()
        os.fsync(stream.fileno())


def _seal(args: argparse.Namespace) -> int:
    job = validate_camera_compute_job(args.job_root)
    selection = _load_selection(args.selection)
    indices = [frame.frame_index for frame in selection]
    with tempfile.TemporaryDirectory(prefix="missioncore-e2-") as temporary:
        temporary_root = Path(temporary)
        stream_path = temporary_root / "camera.mp4"
        frames_root = temporary_root / "frames"
        timeline_path = temporary_root / "timeline.jsonl"
        _reconstruct_stream(job.job_root, stream_path)
        timestamps = _timeline(stream_path, job.timeline_start_seconds, job.segment_count)
        _extract_frames(stream_path, indices, frames_root)
        _write_timeline(timeline_path, timestamps)
        pack = prepare_recorded_evaluation_pack(
            job_root=job.job_root,
            qualification_root=args.qualification_root,
            valid_fov_root=args.valid_fov_root,
            decoded_frames_root=frames_root,
            timeline_path=timeline_path,
            output_root=args.output_root,
            selection=selection,
            decoder_version=_ffmpeg_version(),
            selection_document_sha256=_sha256(args.selection.expanduser().resolve(strict=True)),
            producer_files=(
                (
                    "evaluation_pack.py",
                    _sha256(Path(str(evaluation_pack_module.__file__)).resolve(strict=True)),
                ),
                (
                    "prepare_e2_evaluation_pack.py",
                    _sha256(Path(__file__).resolve(strict=True)),
                ),
            ),
        )
    print(
        json.dumps(
            {
                "schema_version": "missioncore.e2-evaluation-preparation/v1",
                "generation_id": pack.generation_id,
                "root": str(pack.root),
                "job_id": pack.job_id,
                "input_sha256": pack.input_sha256,
                "qualification_generation_id": pack.qualification_generation_id,
                "valid_fov_generation_id": pack.valid_fov_generation_id,
                "calibration_sha256": pack.calibration_sha256,
                "frame_count": len(pack.frames),
                "anchor_count": sum(frame.role == "anchor" for frame in pack.frames),
                "temporal_frame_count": sum(frame.role == "temporal" for frame in pack.frames),
                "annotation_state": "unannotated",
            },
            ensure_ascii=False,
            indent=2,
            sort_keys=True,
        )
    )
    return 0


def main() -> int:
    args = _arguments()
    if args.command == "candidates":
        return _candidates(args)
    if args.command == "seal":
        return _seal(args)
    raise RuntimeError("unknown E2 command")


if __name__ == "__main__":
    raise SystemExit(main())
