#!/usr/bin/env python3
"""Render a long video in chunks and join them without re-encoding, resumably.

An hour-long timeline in one ffmpeg run is a bad bet: a single failure at minute 52
throws away everything, memory grows with the filter graph, and there is no way to
inspect the middle without waiting for the end. Chunking fixes all three.

What makes the join seamless:

  * every chunk is encoded with the *same* settings and forced to open on a
    keyframe, which is what the concat demuxer needs to copy streams through;
  * chunk boundaries are snapped to the frame grid, so no chunk is a fractional
    frame long;
  * audio is optionally re-encoded in the join pass only. AAC carries encoder
    priming at the start of every chunk, and copying it through can leave a click
    at each seam — re-encoding just the audio costs seconds and removes it while
    the video is still copied.

Resume is by content, not by filename: each chunk records the exact parameters it
was rendered with, and a chunk is only reused when those match and the file still
probes as valid.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import subprocess
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from media_probe import probe  # noqa: E402
from render_core import PROFILES, capabilities, decode_arguments, output_arguments  # noqa: E402
from skill_config import emit, require_binary, utf8_stdout  # noqa: E402


def ff(value: float) -> str:
    return f"{value:.6f}".rstrip("0").rstrip(".") or "0"


def signature(payload: dict) -> str:
    return hashlib.sha256(
        json.dumps(payload, sort_keys=True, ensure_ascii=False).encode()
    ).hexdigest()[:16]


def main() -> None:
    utf8_stdout()
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("input")
    parser.add_argument("output")
    parser.add_argument("--chunk-seconds", type=float, default=120.0)
    parser.add_argument("--work-dir", help="Where chunks live (default: OUTPUT.chunks)")
    parser.add_argument("--vf", help="Video filter chain applied to every chunk")
    parser.add_argument("--af", help="Audio filter chain applied to every chunk")
    parser.add_argument("--profile", default="delivery", choices=list(PROFILES))
    parser.add_argument("--codec", default="h264", choices=["h264", "hevc"])
    parser.add_argument("--cpu", action="store_true")
    parser.add_argument("--hwaccel", action="store_true")
    parser.add_argument("--start", type=float, default=0.0)
    parser.add_argument("--duration", type=float, help="Limit total length")
    parser.add_argument("--chapters", help="ffmetadata file to embed as chapters")
    parser.add_argument("--audio-copy", action="store_true",
                        help="Copy audio in the join pass instead of re-encoding it")
    parser.add_argument("--fresh", action="store_true", help="Ignore existing chunks")
    parser.add_argument("--keep-chunks", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    require_binary("ffmpeg")
    source = Path(args.input).resolve()
    if not source.is_file():
        raise SystemExit(f"Not found: {source}")
    info = probe(source)
    total = float(info.get("duration") or 0.0)
    if total <= 0:
        raise SystemExit(f"Could not determine the duration of {source}")
    fps = float((info.get("video") or {}).get("avg_frame_rate")
                or (info.get("video") or {}).get("r_frame_rate") or 30)
    span = min(total - args.start, args.duration) if args.duration else total - args.start
    if span <= 0:
        raise SystemExit("Nothing to render: --start is past the end of the source")

    # Snap the chunk length to whole frames so no chunk is a fractional frame long.
    chunk_frames = max(1, int(round(args.chunk_seconds * fps)))
    chunk_length = chunk_frames / fps
    count = max(1, int(-(-span // chunk_length)))

    work = Path(args.work_dir) if args.work_dir else Path(str(args.output) + ".chunks")
    work.mkdir(parents=True, exist_ok=True)
    caps = capabilities()
    encoder_arguments, selection = output_arguments(
        args.codec, args.profile, prefer_gpu=not args.cpu, caps=caps
    )

    common = {
        "source": str(source),
        "source_bytes": info.get("bytes"),
        "vf": args.vf, "af": args.af,
        "encoder": selection["encoder"], "profile": args.profile,
        "codec": args.codec, "fps": fps,
        "chunk_frames": chunk_frames,
    }
    plan = []
    for index in range(count):
        start = args.start + index * chunk_length
        length = min(chunk_length, args.start + span - start)
        plan.append({
            "index": index,
            "start": round(start, 6),
            "length": round(length, 6),
            "frames": int(round(length * fps)),
            "path": str(work / f"chunk_{index:04d}.mp4"),
            "signature": signature({**common, "start": round(start, 6),
                                    "length": round(length, 6)}),
        })

    report = {
        "source": str(source),
        "total_duration": round(total, 3),
        "rendering": round(span, 3),
        "fps": fps,
        "chunk_seconds": round(chunk_length, 4),
        "chunks": count,
        "encoder": selection,
        "work_dir": str(work.resolve()),
    }
    if args.dry_run:
        report["plan"] = plan
        emit(report)
        return

    state_path = work / "chunks.json"
    previous: dict[str, dict] = {}
    if state_path.exists() and not args.fresh:
        try:
            for item in json.loads(state_path.read_text(encoding="utf-8")).get("chunks", []):
                previous[item["signature"]] = item
        except json.JSONDecodeError:
            previous = {}

    rendered = []
    reused = 0
    started_at = time.perf_counter()
    for item in plan:
        path = Path(item["path"])
        cached = previous.get(item["signature"])
        if cached and path.is_file() and not args.fresh:
            check = probe(path)
            if check.get("duration") and abs(check["duration"] - item["length"]) < 0.5:
                rendered.append({**item, "reused": True,
                                 "actual_duration": check["duration"]})
                reused += 1
                continue
        command = [
            "ffmpeg", "-y", "-v", "error", "-nostdin",
            *decode_arguments(args.hwaccel, caps),
            "-ss", ff(item["start"]), "-t", ff(item["length"]),
            "-i", str(source),
        ]
        if args.vf:
            command += ["-vf", args.vf]
        if args.af:
            command += ["-af", args.af]
        command += [
            *encoder_arguments,
            # Open every chunk on a keyframe: the concat demuxer copies streams and
            # cannot start a chunk mid-GOP.
            "-force_key_frames", "expr:eq(n,0)",
            "-reset_timestamps", "1",
            str(path),
        ]
        chunk_started = time.perf_counter()
        completed = subprocess.run(command, capture_output=True, text=True, errors="replace")
        if completed.returncode:
            report.update({
                "status": "chunk_failed",
                "failed_chunk": item["index"],
                "rendered_so_far": len(rendered),
                "tail": (completed.stderr or "").strip()[-1200:],
                "resume_hint": (
                    "Запусти ту же команду снова — уже готовые чанки будут "
                    "переиспользованы, рендер продолжится с упавшего."
                ),
            })
            state_path.write_text(
                json.dumps({"common": common, "chunks": rendered}, ensure_ascii=False,
                           indent=2), encoding="utf-8")
            emit(report)
            raise SystemExit(completed.returncode)
        check = probe(path)
        rendered.append({
            **item, "reused": False,
            "actual_duration": check.get("duration"),
            "seconds_to_render": round(time.perf_counter() - chunk_started, 2),
        })
        state_path.write_text(
            json.dumps({"common": common, "chunks": rendered}, ensure_ascii=False, indent=2),
            encoding="utf-8")

    listing = work / "concat.txt"
    # Absolute paths: the concat demuxer resolves relative entries against the
    # directory of the list file, not the working directory, so a relative path
    # silently becomes work_dir/work_dir/chunk.mp4.
    listing.write_text(
        "\n".join(f"file '{Path(item['path']).resolve().as_posix()}'" for item in rendered)
        + "\n",
        encoding="utf-8",
    )

    output = Path(args.output).resolve()
    output.parent.mkdir(parents=True, exist_ok=True)
    join = ["ffmpeg", "-y", "-v", "error", "-nostdin",
            "-f", "concat", "-safe", "0", "-i", str(listing)]
    if args.chapters:
        join += ["-i", str(Path(args.chapters).resolve()), "-map_metadata", "1"]
    join += ["-map", "0:v", "-map", "0:a?", "-c:v", "copy"]
    if args.audio_copy:
        join += ["-c:a", "copy"]
    else:
        # Re-encoding only the audio removes the click that AAC priming leaves at
        # each seam, and costs a fraction of the time a video re-encode would.
        join += ["-c:a", "aac", "-b:a", "192k", "-ar", "48000", "-ac", "2"]
    join += ["-movflags", "+faststart", str(output)]
    joined = subprocess.run(join, capture_output=True, text=True, errors="replace")
    if joined.returncode:
        report.update({"status": "join_failed",
                       "tail": (joined.stderr or "").strip()[-1200:],
                       "concat_list": str(listing)})
        emit(report)
        raise SystemExit(joined.returncode)

    final = probe(output)
    expected = sum(item["length"] for item in plan)
    actual = float(final.get("duration") or 0.0)
    drift = round(actual - expected, 3)
    tolerance = max(0.2, 2.0 / fps * count ** 0.5)

    if not args.keep_chunks:
        for item in rendered:
            Path(item["path"]).unlink(missing_ok=True)
        listing.unlink(missing_ok=True)

    report.update({
        "status": "ok" if abs(drift) <= tolerance else "duration_mismatch",
        "output": str(output),
        "expected_duration": round(expected, 3),
        "actual_duration": actual,
        "duration_drift": drift,
        "tolerance": round(tolerance, 3),
        "chunks_rendered": sum(1 for item in rendered if not item["reused"]),
        "chunks_reused": reused,
        "wall_seconds": round(time.perf_counter() - started_at, 1),
        "realtime_factor": round(span / max(1e-6, time.perf_counter() - started_at), 2),
        "audio": "copied" if args.audio_copy else "re-encoded in the join pass",
        "chapters_embedded": bool(args.chapters),
        "chunks_kept": args.keep_chunks,
        "slowest_chunk": max(
            (item for item in rendered if not item["reused"]),
            key=lambda item: item.get("seconds_to_render") or 0,
            default=None,
        ),
    })
    emit(report)
    if abs(drift) > tolerance:
        raise SystemExit(2)


if __name__ == "__main__":
    main()