from __future__ import annotations

import json
import subprocess
from pathlib import Path
from typing import Any


def extract_hdr10plus_metadata(
    *,
    ffmpeg: Path,
    hdr10plus_tool: Path,
    source: Path,
    output: Path,
) -> None:
    """Extract HDR10+ JSON without creating a full intermediate video file."""
    output.parent.mkdir(parents=True, exist_ok=True)
    ffmpeg_process = subprocess.Popen(
        [
            str(ffmpeg),
            "-v", "error", "-i", str(source), "-map", "0:v:0",
            "-c", "copy", "-bsf:v", "hevc_mp4toannexb", "-f", "hevc", "-",
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    assert ffmpeg_process.stdout is not None
    try:
        extracted = subprocess.run(
            [str(hdr10plus_tool), "extract", "-o", str(output), "-"],
            stdin=ffmpeg_process.stdout,
            capture_output=True,
            text=True,
        )
    finally:
        ffmpeg_process.stdout.close()
    ffmpeg_stderr = ffmpeg_process.stderr.read().decode(errors="replace")
    ffmpeg_process.stderr.close()
    ffmpeg_returncode = ffmpeg_process.wait()
    if ffmpeg_returncode != 0:
        raise RuntimeError(
            f"FFmpeg could not extract the HEVC stream from {source}:\n{ffmpeg_stderr.strip()}"
        )
    if extracted.returncode != 0:
        raise RuntimeError(
            f"hdr10plus_tool could not extract metadata from {source}:\n"
            f"{extracted.stderr.strip() or extracted.stdout.strip()}"
        )
    if not output.is_file():
        raise RuntimeError(f"hdr10plus_tool did not create metadata JSON for {source}")


def select_hdr10plus_frames(
    *,
    source_json: Path,
    source_frames: list[int],
    output_json: Path,
) -> None:
    """Keep metadata for the encoded source frames and rebuild scene indices."""
    payload = json.loads(source_json.read_text(encoding="utf-8"))
    entries = payload.get("SceneInfo")
    if not isinstance(entries, list) or not entries:
        raise RuntimeError(f"HDR10+ JSON has no SceneInfo entries: {source_json}")
    if not source_frames:
        raise RuntimeError("Cannot create HDR10+ metadata for an empty output")
    if min(source_frames) < 0 or max(source_frames) >= len(entries):
        raise RuntimeError(
            "HDR10+ metadata frame count does not cover the validated source frames "
            f"({len(entries)} metadata entries, highest source frame {max(source_frames)})"
        )

    selected = [dict(entries[index]) for index in source_frames]
    _reindex_scenes(selected)
    payload["SceneInfo"] = selected
    payload["SceneInfoSummary"] = _scene_summary(selected)
    output_json.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")


def _reindex_scenes(entries: list[dict[str, Any]]) -> None:
    scene_id = 0
    scene_frame_index = 0
    previous: dict[str, Any] | None = None
    for sequence_frame_index, entry in enumerate(entries):
        if previous is not None and _scene_payload(entry) != _scene_payload(previous):
            scene_id += 1
            scene_frame_index = 0
        entry["SceneFrameIndex"] = scene_frame_index
        entry["SceneId"] = scene_id
        entry["SequenceFrameIndex"] = sequence_frame_index
        scene_frame_index += 1
        previous = entry


def _scene_payload(entry: dict[str, Any]) -> dict[str, Any]:
    return {
        key: value
        for key, value in entry.items()
        if key not in {"SceneFrameIndex", "SceneId", "SequenceFrameIndex"}
    }


def _scene_summary(entries: list[dict[str, Any]]) -> dict[str, list[int]]:
    starts = [
        index
        for index, entry in enumerate(entries)
        if entry.get("SceneFrameIndex") == 0
    ]
    return {
        "SceneFirstFrameIndex": starts,
        "SceneFrameNumbers": [
            next_start - start
            for start, next_start in zip(starts, [*starts[1:], len(entries)])
        ],
    }