#!/usr/bin/env python3
"""Render a recut plan: trim, retime, and keep picture and sound on one time map.

Video uses `setpts=PTS/speed`, audio uses `rubberband` at the same factor, so a
segment played at 1.06x has its picture and its voice stretched identically and
lip sync is preserved by construction. `rubberband` is chosen over `atempo`
because it preserves formants — a voice sped up with `atempo` beyond a few percent
starts to sound thin.

The render is verified after the fact: the output duration is compared with the
plan's arithmetic, and a mismatch fails the step rather than being reported as
success.
"""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
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 canvas as resolve_canvas, emit, require_binary, utf8_stdout  # noqa: E402


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


def audio_retime(speed: float, engine: str) -> str:
    """Audio time-stretch at `speed`, pitch preserved."""
    if abs(speed - 1.0) < 1e-4:
        return ""
    if engine == "rubberband":
        return f",rubberband=tempo={speed:.6f}:pitchq=quality:transients=smooth"
    # atempo is limited to 0.5-2.0 per instance; chain it if ever asked for more.
    remaining = speed
    stages = []
    while remaining > 2.0:
        stages.append("atempo=2.0")
        remaining /= 2.0
    while remaining < 0.5:
        stages.append("atempo=0.5")
        remaining /= 0.5
    stages.append(f"atempo={remaining:.6f}")
    return "," + ",".join(stages)


def main() -> None:
    utf8_stdout()
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("plan", help="recut_plan.json from plan_speech_recut.py")
    parser.add_argument("output")
    parser.add_argument("--canvas", default="reels",
                        help="Target canvas; use 'source' to keep the original geometry")
    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("--audio-engine", default="rubberband",
                        choices=["rubberband", "atempo"])
    parser.add_argument("--grade", help="Extra ffmpeg video filter applied to every segment")
    parser.add_argument("--no-crop", action="store_true",
                        help="Keep the source geometry instead of fitting the canvas")
    parser.add_argument("--filter-report", help="Write the filter graph here")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    require_binary("ffmpeg")
    plan = json.loads(Path(args.plan).read_text(encoding="utf-8"))
    segments = plan.get("segments") or []
    if not segments:
        raise SystemExit("Plan has no segments")
    source = Path(plan["source"])
    if not source.is_file():
        raise SystemExit(f"Source from the plan is missing: {source}")
    info = probe(source)

    if args.canvas == "source" or args.no_crop:
        width = (info.get("video") or {}).get("display_width") or 1080
        height = (info.get("video") or {}).get("display_height") or 1920
        fps = int(round((info.get("video") or {}).get("avg_frame_rate") or 30))
        geometry = f"setsar=1,fps={fps}"
    else:
        board = resolve_canvas(args.canvas)
        width, height, fps = board.width, board.height, board.fps
        geometry = (
            f"scale={width}:{height}:force_original_aspect_ratio=increase:flags=lanczos,"
            f"crop={width}:{height},setsar=1,fps={fps}"
        )

    expected = sum(
        (item["source_out"] - item["source_in"]) / item.get("speed", 1.0)
        for item in segments
    )
    report = {
        "plan": str(Path(args.plan).resolve()),
        "source": str(source),
        "segments": len(segments),
        "canvas": f"{width}x{height}@{fps}",
        "speeds": sorted({item.get("speed", 1.0) for item in segments}),
        "expected_duration": round(expected, 4),
        "audio_engine": args.audio_engine,
    }
    if args.dry_run:
        emit(report)
        return

    graph: list[str] = []
    video_labels: list[str] = []
    audio_labels: list[str] = []
    for index, item in enumerate(segments):
        begin = float(item["source_in"])
        finish = float(item["source_out"])
        speed = float(item.get("speed", 1.0))
        chain = (
            f"[0:v]trim=start={ff(begin)}:end={ff(finish)},setpts=PTS-STARTPTS"
        )
        if abs(speed - 1.0) > 1e-4:
            chain += f",setpts=PTS/{speed:.6f}"
        chain += f",{geometry}"
        if args.grade:
            chain += f",{args.grade}"
        chain += f",format=yuv420p[v{index}]"
        graph.append(chain)
        video_labels.append(f"[v{index}]")

        audio_chain = (
            f"[0:a]atrim=start={ff(begin)}:end={ff(finish)},asetpts=PTS-STARTPTS"
            f"{audio_retime(speed, args.audio_engine)}"
            f",aformat=sample_rates=48000:channel_layouts=stereo[a{index}]"
        )
        graph.append(audio_chain)
        audio_labels.append(f"[a{index}]")

    interleaved = "".join(f"{v}{a}" for v, a in zip(video_labels, audio_labels))
    graph.append(f"{interleaved}concat=n={len(segments)}:v=1:a=1[vout][aout]")

    filter_graph = ";\n".join(graph)
    report_path = (
        Path(args.filter_report) if args.filter_report
        else Path(args.output).with_suffix(".filter.txt")
    )
    report_path.parent.mkdir(parents=True, exist_ok=True)
    report_path.write_text(filter_graph + "\n", encoding="utf-8")

    caps = capabilities()
    encoder_arguments, selection = output_arguments(
        args.codec, args.profile, prefer_gpu=not args.cpu, caps=caps
    )
    output = Path(args.output).resolve()
    output.parent.mkdir(parents=True, exist_ok=True)
    command = [
        "ffmpeg", "-y", "-v", "warning", "-nostdin",
        *decode_arguments(args.hwaccel, caps),
        "-i", str(source),
        "-filter_complex_script", str(report_path),
        "-map", "[vout]", "-map", "[aout]",
        *encoder_arguments,
        str(output),
    ]
    completed = subprocess.run(command, capture_output=True, text=True, errors="replace")
    if completed.returncode:
        report.update({"status": "render_failed",
                       "tail": (completed.stderr or "").strip()[-1500:],
                       "filter_report": str(report_path)})
        emit(report)
        raise SystemExit(completed.returncode)

    actual = probe(output)
    drift = round((actual.get("duration") or 0.0) - expected, 4)
    tolerance = max(0.12, 3.0 / fps)
    report.update({
        "status": "ok" if abs(drift) <= tolerance else "duration_mismatch",
        "output": str(output),
        "actual_duration": actual.get("duration"),
        "duration_drift": drift,
        "tolerance": round(tolerance, 4),
        "encoder": selection,
        "filter_report": str(report_path),
        "note": (
            "Скорость менялась только на склейках, видео и звук растянуты одним "
            "множителем — липсинк сохранён."
        ),
    })
    emit(report)
    if abs(drift) > tolerance:
        raise SystemExit(2)


if __name__ == "__main__":
    main()