#!/usr/bin/env python3
"""Join clips with real transitions in one ffmpeg pass, and place the SFX for them.

The important arithmetic, because getting it wrong is the usual cause of a video
that drifts out of sync with its captions: `xfade` *overlaps* the two inputs, so
every transition shortens the result by its own duration. A sequence of clips with
durations d0..dn and transitions t1..tn ends at `sum(d) - sum(t)`. This script
reports that number, writes the shifted timeline of every cut point, and refuses
to run when a transition is longer than the clip it has to overlap.

Audio is crossfaded over the same windows with `acrossfade`, so picture and sound
share one time map rather than being cut twice.
"""
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 skill_config import canvas as resolve_canvas, emit, require_binary, utf8_stdout  # noqa: E402
from transition_lib import needs_full_chroma, recipe, xfade_arguments  # noqa: E402


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


def load_plan(path: Path) -> dict:
    plan = json.loads(path.read_text(encoding="utf-8"))
    if "clips" not in plan or len(plan["clips"]) < 1:
        raise SystemExit("Plan needs a non-empty 'clips' list")
    return plan


def normalise_chain(width: int, height: int, fps: int, grade: str | None) -> str:
    chain = (
        f"scale={width}:{height}:force_original_aspect_ratio=increase:flags=lanczos,"
        f"crop={width}:{height},setsar=1,fps={fps},format=yuv420p"
    )
    if grade:
        chain += f",{grade}"
    return chain


def main() -> None:
    utf8_stdout()
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("plan", help="transition_plan.json")
    parser.add_argument("output")
    parser.add_argument("--canvas", default="reels")
    parser.add_argument("--fps", type=int)
    parser.add_argument("--crf", type=int, default=17)
    parser.add_argument("--preset", default="slow")
    parser.add_argument("--grade", help="Extra ffmpeg filter applied to every clip")
    parser.add_argument("--filter-report", help="Write the filter graph here for audit")
    parser.add_argument("--dry-run", action="store_true", help="Report timing only")
    args = parser.parse_args()

    require_binary("ffmpeg")
    plan = load_plan(Path(args.plan))
    board = resolve_canvas(plan.get("canvas", args.canvas))
    width = int(plan.get("width", board.width))
    height = int(plan.get("height", board.height))
    fps = int(args.fps or plan.get("fps", board.fps))

    clips = plan["clips"]
    resolved: list[dict] = []
    for index, clip in enumerate(clips):
        source = Path(clip["path"]).resolve()
        if not source.is_file():
            raise SystemExit(f"Clip {index} not found: {source}")
        info = probe(source)
        start = float(clip.get("in", 0.0))
        available = info.get("duration") or 0.0
        end = float(clip["out"]) if clip.get("out") is not None else available
        if end <= start:
            raise SystemExit(f"Clip {index} has a non-positive length: in={start} out={end}")
        resolved.append({
            "index": index,
            "path": source,
            "in": start,
            "out": end,
            "duration": end - start,
            "has_audio": bool(info.get("has_audio")),
            "kind": info.get("kind"),
            "still": info.get("kind") == "image",
        })

    # Transitions sit between clips, so there is always one fewer than the clips.
    incoming = plan.get("transitions") or []
    if len(incoming) not in (0, len(resolved) - 1):
        raise SystemExit(
            f"{len(resolved)} clips need {len(resolved) - 1} transitions, got {len(incoming)}"
        )
    if not incoming:
        incoming = [{"id": "cut"} for _ in range(max(0, len(resolved) - 1))]

    steps: list[dict] = []
    problems: list[str] = []
    for position, entry in enumerate(incoming):
        item = recipe(entry.get("id", "cut"))
        frames = entry.get("frames")
        duration = 0.0 if item.is_cut else item.duration(fps, frames)
        before = resolved[position]["duration"]
        after = resolved[position + 1]["duration"]
        # An xfade cannot consume more material than either side actually has.
        headroom = min(before, after) * 0.6
        if duration > headroom:
            problems.append(
                f"{item.id} between clip {position} and {position + 1}: "
                f"{duration:.2f}s exceeds 60% of the shorter clip ({headroom:.2f}s); "
                f"shortened"
            )
            duration = max(0.0, headroom)
        steps.append({"position": position, "recipe": item, "duration": duration,
                      "frames": round(duration * fps)})

    total_source = sum(clip["duration"] for clip in resolved)
    total_overlap = sum(step["duration"] for step in steps)
    expected = total_source - total_overlap

    # Where each cut lands on the output timeline, after every earlier overlap.
    cut_points: list[dict] = []
    cursor = 0.0
    consumed = 0.0
    for position, clip in enumerate(resolved[:-1]):
        cursor += clip["duration"]
        step = steps[position]
        cut_at = cursor - consumed - step["duration"] / 2
        cut_points.append({
            "index": position,
            "transition": step["recipe"].id,
            "duration": round(step["duration"], 4),
            "frames": step["frames"],
            "start": round(cut_at - step["duration"] / 2, 4),
            "centre": round(cut_at, 4),
            "end": round(cut_at + step["duration"] / 2, 4),
        })
        consumed += step["duration"]

    sfx_events: list[dict] = []
    for point, step in zip(cut_points, steps):
        for cue in step["recipe"].sfx:
            offset = float(cue.get("offset_frames", 0)) / fps
            sfx_events.append({
                "sound": cue["id"],
                "at": round(max(0.0, point["centre"] + offset), 4),
                "gain_db": float(cue.get("gain_db", -18)),
                "pan": cue.get("pan"),
                "transition": step["recipe"].id,
                "cut_index": point["index"],
            })
    sfx_events.sort(key=lambda item: item["at"])

    report = {
        "status": "ok" if not problems else "ok_with_adjustments",
        "output": str(Path(args.output).resolve()),
        "canvas": f"{width}x{height}@{fps}",
        "clips": len(resolved),
        "transitions": [
            {"id": step["recipe"].id, "family": step["recipe"].family,
             "frames": step["frames"], "seconds": round(step["duration"], 4)}
            for step in steps
        ],
        "source_duration": round(total_source, 4),
        "overlap_total": round(total_overlap, 4),
        "expected_duration": round(expected, 4),
        "duration_note": (
            "Каждый переход перекрывает соседние клипы, поэтому итог короче суммы "
            "исходников ровно на overlap_total. Титры и музыка обязаны читать "
            "cut_points, а не исходные времена."
        ),
        "cut_points": cut_points,
        "sfx_events": sfx_events,
        "adjustments": problems,
    }

    if args.dry_run:
        emit(report)
        return

    # ---------------------------------------------------------------- filtergraph
    command = ["ffmpeg", "-y", "-v", "warning"]
    # Set below if any recipe in the plan uses a custom expression. `xfade` slices
    # the frame across filter threads, and every slice evaluates the *same*
    # AVExpr, so the `st()` / `ld()` register file is shared. A recipe that stores
    # a per-pixel value — a band index, a displaced X — has that register
    # overwritten by another thread between the store and the load, and the frame
    # comes out as coloured noise instead of the effect. Recipes whose `st()`
    # slots depend only on Q survive, which is why the failure looked like "some
    # transitions are broken" rather than a systematic bug.
    single_thread = False
    for clip in resolved:
        if clip["still"]:
            command += ["-loop", "1", "-t", ff(clip["duration"]), "-i", str(clip["path"])]
        else:
            command += ["-ss", ff(clip["in"]), "-to", ff(clip["out"]), "-i", str(clip["path"])]

    graph: list[str] = []
    base_chain = normalise_chain(width, height, fps, args.grade)
    for clip in resolved:
        index = clip["index"]
        chain = base_chain
        if not clip["still"]:
            chain = f"setpts=PTS-STARTPTS,{chain}"
        graph.append(f"[{index}:v]{chain}[v{index}]")
        if clip["has_audio"]:
            graph.append(
                f"[{index}:a]asetpts=PTS-STARTPTS,"
                f"aformat=sample_rates=48000:channel_layouts=stereo[a{index}]"
            )
        else:
            # Silence keeps the audio graph symmetric; without it acrossfade would
            # have nothing to fade and the streams would go out of step.
            graph.append(
                f"anullsrc=r=48000:cl=stereo,atrim=duration={ff(clip['duration'])},"
                f"asetpts=PTS-STARTPTS[a{index}]"
            )

    video_label = "v0"
    audio_label = "a0"
    running = resolved[0]["duration"]
    for position, step in enumerate(steps):
        item = step["recipe"]
        duration = step["duration"]
        next_video = f"v{position + 1}"
        next_audio = f"a{position + 1}"
        out_video = f"xv{position}"
        out_audio = f"xa{position}"
        if item.is_cut or duration <= 0:
            graph.append(f"[{video_label}][{next_video}]concat=n=2:v=1:a=0[{out_video}]")
            graph.append(f"[{audio_label}][{next_audio}]concat=n=2:v=0:a=1[{out_audio}]")
            running += resolved[position + 1]["duration"]
        else:
            pre = item.ffmpeg.get("pre_out")
            if pre:
                graph.append(f"[{video_label}]{pre}[pv{position}]")
                video_label = f"pv{position}"
            pre_in = item.ffmpeg.get("pre_in")
            if pre_in:
                graph.append(f"[{next_video}]{pre_in}[pn{position}]")
                next_video = f"pn{position}"
            full_chroma = needs_full_chroma(item.ffmpeg)
            if full_chroma:
                single_thread = True
                # A custom expression addresses pixels as W x H on every plane,
                # so both sides go to 4:4:4 for the transition and come back after.
                graph.append(f"[{video_label}]format=yuv444p[cf{position}a]")
                graph.append(f"[{next_video}]format=yuv444p[cf{position}b]")
                video_label = f"cf{position}a"
                next_video = f"cf{position}b"
            offset = running - duration
            target = f"{out_video}_raw" if full_chroma else out_video
            graph.append(
                f"[{video_label}][{next_video}]"
                f"{xfade_arguments(item.ffmpeg, duration, offset)}[{target}]"
            )
            if full_chroma:
                graph.append(f"[{target}]format=yuv420p[{out_video}]")
            graph.append(
                f"[{audio_label}][{next_audio}]"
                f"acrossfade=d={ff(duration)}:c1=tri:c2=tri[{out_audio}]"
            )
            running += resolved[position + 1]["duration"] - duration
        video_label = out_video
        audio_label = out_audio

    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")

    output = Path(args.output).resolve()
    output.parent.mkdir(parents=True, exist_ok=True)
    if single_thread:
        # Global option, so it is inserted ahead of the inputs rather than
        # appended: ffmpeg reads it once for the whole filtergraph.
        command[1:1] = ["-filter_complex_threads", "1"]
        report["filter_complex_threads"] = 1
        report["thread_note"] = (
            "План содержит custom-выражение. Графа считается в один поток: "
            "срезы кадра делят один набор регистров st()/ld(), и попиксельное "
            "значение в них затирается соседним потоком — кадр выходит шумом."
        )
    command += [
        "-filter_complex_script", str(report_path),
        "-map", f"[{video_label}]",
        "-map", f"[{audio_label}]",
        "-c:v", "libx264", "-preset", args.preset, "-crf", str(args.crf),
        "-profile:v", "high", "-pix_fmt", "yuv420p", "-r", str(fps),
        "-c:a", "aac", "-b:a", "192k", "-ar", "48000", "-ac", "2",
        "-movflags", "+faststart",
        str(output),
    ]
    completed = subprocess.run(command)
    if completed.returncode:
        report["status"] = "render_failed"
        report["filter_report"] = str(report_path)
        emit(report)
        raise SystemExit(completed.returncode)

    actual = probe(output)
    report["filter_report"] = str(report_path)
    report["actual_duration"] = actual.get("duration")
    drift = None
    if actual.get("duration"):
        drift = round(actual["duration"] - expected, 4)
    report["duration_drift"] = drift
    report["drift_ok"] = drift is not None and abs(drift) <= max(0.08, 2.0 / fps)
    if not report["drift_ok"]:
        report["status"] = "duration_mismatch"
    plan_out = Path(args.output).with_suffix(".transitions.json")
    plan_out.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    report["report"] = str(plan_out)
    emit(report)
    if not report["drift_ok"]:
        raise SystemExit(2)


if __name__ == "__main__":
    main()