#!/usr/bin/env python3
"""Assemble six local H3 shots into a 30-second, 1344x768/24fps MP4.

Requires Python 3.10+, ffmpeg and ffprobe on PATH. No network access.
Inputs: shot01.mp4 ... shot06.mp4, each containing >=5s of video and audio.
Optional narration is a locally recorded >=30s audio track. In that mode,
generate H3 shots without dialogue: their audio is retained as quiet ambience.
Optional local music and UTF-8 SRT subtitles are supported. Narration ducks
the ambience/music bed automatically. A JSON verification report accompanies
the final MP4. Input files and existing outputs are never replaced.
"""
from __future__ import annotations

import argparse
from fractions import Fraction
import json
import math
from pathlib import Path
import re
import shutil
import subprocess
import sys
import tempfile


def run(args: list[str]) -> subprocess.CompletedProcess[str]:
    result = subprocess.run(args, capture_output=True, text=True, encoding="utf-8", errors="replace")
    if result.returncode:
        raise RuntimeError(f"Command failed: {args[0]}\n{result.stderr[-4500:]}")
    return result


def probe(path: Path, count_frames: bool = False) -> dict:
    args = ["ffprobe", "-v", "error"]
    if count_frames:
        args += ["-count_frames"]
    args += ["-show_streams", "-show_format", "-of", "json", str(path)]
    return json.loads(run(args).stdout)


def duration(stream: dict) -> float:
    # H3/MP4 stream durations are required; a container duration alone can hide
    # an audio track that ends before the picture.
    value = stream.get("duration")
    try:
        return float(value)
    except (TypeError, ValueError):
        ticks, base = stream.get("duration_ts"), stream.get("time_base", "")
        if ticks is not None and "/" in base:
            numerator, denominator = map(int, base.split("/"))
            return float(ticks) * numerator / denominator
    raise ValueError("Stream has no reliable duration; export this input as MP4 or WAV first.")


def require_media(path: Path, seconds: float, need_video: bool) -> dict:
    if not path.is_file():
        raise ValueError(f"Missing input: {path}")
    data = probe(path)
    kinds = ("video", "audio") if need_video else ("audio",)
    for kind in kinds:
        streams = [s for s in data["streams"] if s.get("codec_type") == kind]
        if not streams:
            raise ValueError(f"{path.name}: missing {kind} track.")
        length = duration(streams[0])
        if length + 0.0001 < seconds:
            raise ValueError(f"{path.name}: {kind} is {length:.3f}s; need at least {seconds:.3f}s. Regenerate/export a longer shot.")
    return data


def loudness_filter(inputs: list[str], mix: str) -> str:
    analysis = run(["ffmpeg", "-hide_banner", "-nostats", *inputs,
                    "-filter_complex", mix + ";[mix]loudnorm=I=-16:TP=-2:LRA=7:print_format=json[analysis]",
                    "-map", "[analysis]", "-vn", "-f", "null", "-"]).stderr
    blocks = re.findall(r'\{\s*"input_i"[\s\S]*?\}', analysis)
    if not blocks:
        raise RuntimeError("FFmpeg did not return loudness measurements.")
    stats = json.loads(blocks[-1])
    values = [float(stats[k]) for k in ("input_i", "input_tp", "input_lra", "input_thresh", "target_offset")]
    if not all(math.isfinite(value) for value in values):
        # Silence has no finite integrated loudness; keep the track silent.
        return "anull"
    return ("loudnorm=I=-16:TP=-2:LRA=7:linear=true:"
            f"measured_I={stats['input_i']}:measured_TP={stats['input_tp']}:"
            f"measured_LRA={stats['input_lra']}:measured_thresh={stats['input_thresh']}:"
            f"offset={stats['target_offset']}")


def validate_srt(path: Path) -> int:
    if not path.is_file():
        raise ValueError(f"Missing subtitle file: {path}")
    text = path.read_text(encoding="utf-8-sig").replace("\r\n", "\n").strip()
    if not text:
        raise ValueError("SRT file is empty.")
    pattern = re.compile(r"^(\d{2}):(\d{2}):(\d{2}),(\d{3}) --> (\d{2}):(\d{2}):(\d{2}),(\d{3})$")
    previous_end = 0.0
    blocks = re.split(r"\n[ \t]*\n", text)
    for index, block in enumerate(blocks, 1):
        lines = block.splitlines()
        if len(lines) < 3 or not lines[0].strip().isdigit():
            raise ValueError(f"SRT cue {index}: expected number, timestamp line, then subtitle text.")
        match = pattern.fullmatch(lines[1].strip())
        if not match or not any(line.strip() for line in lines[2:]):
            raise ValueError(f"SRT cue {index}: use HH:MM:SS,mmm --> HH:MM:SS,mmm and non-empty text.")
        hh, mm, ss, ms, eh, em, es, ez = map(int, match.groups())
        if max(mm, ss, em, es) > 59:
            raise ValueError(f"SRT cue {index}: minutes and seconds must be below 60.")
        start = hh * 3600 + mm * 60 + ss + ms / 1000
        end = eh * 3600 + em * 60 + es + ez / 1000
        if start < previous_end or start >= end or end > 30:
            raise ValueError(f"SRT cue {index}: use non-overlapping ordered cues with 0 <= start < end <= 30s.")
        previous_end = end
    return len(blocks)


def build_audio_mix(narration_index: int | None, music_index: int | None) -> str:
    filters = []
    ambience_gain = "0.16" if narration_index is not None else "1"
    filters.append(f"[0:a]atrim=0:30,asetpts=PTS-STARTPTS,volume={ambience_gain}[amb]")
    if music_index is not None:
        filters.append(f"[{music_index}:a]atrim=0:30,asetpts=PTS-STARTPTS,aresample=48000,"
                       "aformat=channel_layouts=stereo,volume=0.12,"
                       "afade=t=in:st=0:d=0.5,afade=t=out:st=29.5:d=0.5[music]")
        filters.append("[amb][music]amix=inputs=2:duration=first:dropout_transition=0:normalize=0[bed]")
    else:
        filters.append("[amb]anull[bed]")
    if narration_index is not None:
        filters.append(f"[{narration_index}:a]atrim=0:30,asetpts=PTS-STARTPTS,aresample=48000,"
                       "aformat=channel_layouts=stereo,asplit=2[voice][control]")
        filters.append("[bed][control]sidechaincompress=threshold=0.025:ratio=8:attack=20:release=250:"
                       "detection=rms[ducked]")
        filters.append("[ducked][voice]amix=inputs=2:duration=first:dropout_transition=0:normalize=0[mix]")
    else:
        filters.append("[bed]anull[mix]")
    return ";".join(filters)


def measure_final_loudness(path: Path) -> dict:
    result = run(["ffmpeg", "-hide_banner", "-nostats", "-i", str(path), "-map", "0:a:0",
                  "-af", "loudnorm=I=-16:TP=-2:LRA=7:print_format=json", "-vn", "-f", "null", "-"])
    blocks = re.findall(r'\{\s*"input_i"[\s\S]*?\}', result.stderr)
    if not blocks:
        raise RuntimeError("Final audio loudness measurement failed.")
    stats = json.loads(blocks[-1])
    # Strings preserve FFmpeg's -inf measurements for intentionally silent audio.
    return {"integrated_lufs": stats["input_i"], "true_peak_dbtp": stats["input_tp"], "lra_lu": stats["input_lra"]}


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("shots", type=Path, help="Directory containing shot01.mp4 ... shot06.mp4")
    parser.add_argument("--output", type=Path, help="Final MP4 path; default: SHOTS/final-30s.mp4")
    parser.add_argument("--narration", type=Path, help="Optional >=30s local narration WAV/audio; use no-dialogue H3 shots")
    parser.add_argument("--music", type=Path, help="Optional >=30s local music audio; mixed quietly with 0.5s start/end fades")
    parser.add_argument("--srt", type=Path, help="Optional local UTF-8 SRT, ordered non-overlapping cues within 0..30s; selectable MP4 subtitle track")
    args = parser.parse_args()
    for executable in ("ffmpeg", "ffprobe"):
        if not shutil.which(executable):
            raise ValueError(f"Install {executable} and reopen the terminal so it is on PATH.")
    shots_dir = args.shots.expanduser().resolve()
    clips = [shots_dir / f"shot{i:02}.mp4" for i in range(1, 7)]
    output = (args.output or shots_dir / "final-30s.mp4").expanduser().resolve()
    if output.suffix.lower() != ".mp4":
        raise ValueError("--output must end in .mp4")
    if output.exists():
        raise ValueError(f"Output already exists: {output}. Choose a new --output path.")
    report_path = output.with_suffix(".report.json")
    if report_path.exists():
        raise ValueError(f"Report already exists: {report_path}. Choose a new --output path.")
    if not output.parent.is_dir():
        raise ValueError(f"Output directory does not exist: {output.parent}")
    for clip in clips:
        require_media(clip, 5.0, True)
    narration = args.narration.expanduser().resolve() if args.narration else None
    if narration:
        require_media(narration, 30.0, False)
    music = args.music.expanduser().resolve() if args.music else None
    if music:
        require_media(music, 30.0, False)
    srt = args.srt.expanduser().resolve() if args.srt else None
    subtitle_cues = validate_srt(srt) if srt else 0
    with tempfile.TemporaryDirectory(prefix="h3-edit-", dir=output.parent) as scratch:
        temp = Path(scratch)
        for index, clip in enumerate(clips, 1):
            print(f"[{index}/6] Normalize {clip.name}", flush=True)
            vf = ("fps=24,scale=1344:768:force_original_aspect_ratio=decrease:force_divisible_by=2,"
                  "pad=1344:768:(ow-iw)/2:(oh-ih)/2,setsar=1,"
                  "trim=start_frame=0:end_frame=120,setpts=PTS-STARTPTS")
            af = "atrim=start=0:end=5,asetpts=PTS-STARTPTS,aresample=48000,aformat=channel_layouts=stereo"
            normalized = temp / f"clip{index:02}.mov"
            run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-i", str(clip),
                 "-map", "0:v:0", "-map", "0:a:0", "-vf", vf, "-af", af,
                 "-c:v", "libx264", "-crf", "16", "-preset", "medium", "-pix_fmt", "yuv420p", "-bf", "0",
                 "-c:a", "pcm_s16le", "-ar", "48000", "-ac", "2", str(normalized)])
            video = next(s for s in probe(normalized, True)["streams"] if s["codec_type"] == "video")
            if int(video.get("nb_read_frames", 0)) != 120:
                raise ValueError(f"{clip.name}: decoded output has fewer than 120 frames; regenerate this input.")
        (temp / "shots.txt").write_text("".join(f"file 'clip{i:02}.mov'\n" for i in range(1, 7)), encoding="utf-8")
        inputs = ["-f", "concat", "-safe", "0", "-i", str(temp / "shots.txt")]
        narration_index = music_index = None
        next_index = 1
        if narration:
            narration_index = next_index
            next_index += 1
            inputs += ["-i", str(narration)]
        if music:
            music_index = next_index
            next_index += 1
            inputs += ["-i", str(music)]
        mix = build_audio_mix(narration_index, music_index)
        print("Measure and normalize full-film audio...", flush=True)
        normalizer = loudness_filter(inputs, mix)
        final_temp = temp / "final.mp4"
        subtitle_input = ["-i", str(srt)] if srt else []
        subtitle_maps = ["-map", f"{next_index}:s:0", "-c:s", "mov_text", "-metadata:s:s:0", "language=zho"] if srt else []
        run(["ffmpeg", "-hide_banner", "-loglevel", "error", *inputs, *subtitle_input,
             "-filter_complex", mix + f";[mix]{normalizer},aresample=48000[aout]",
             "-map", "0:v:0", "-map", "[aout]", *subtitle_maps, "-c:v", "copy", "-c:a", "aac",
             "-b:a", "192k", "-ar", "48000", "-ac", "2", "-t", "30", "-movflags", "+faststart", str(final_temp)])
        result = probe(final_temp, True)
        video = next(s for s in result["streams"] if s["codec_type"] == "video")
        audio = next(s for s in result["streams"] if s["codec_type"] == "audio")
        if (int(video.get("nb_read_frames", 0)), video["width"], video["height"]) != (720, 1344, 768):
            raise RuntimeError("Final verification failed: expected 720 frames at 1344x768.")
        if (abs(duration(video) - 30.0) > 0.001 or Fraction(video["avg_frame_rate"]) != 24
                or audio["sample_rate"] != "48000" or audio["channels"] != 2):
            raise RuntimeError(f"Final verification failed: video duration={duration(video):.6f}, audio rate={audio['sample_rate']}.")
        subtitle_streams = [s for s in result["streams"] if s.get("codec_type") == "subtitle"]
        if srt and (len(subtitle_streams) != 1 or subtitle_streams[0].get("codec_name") != "mov_text"):
            raise RuntimeError("Final verification failed: expected one mov_text subtitle track.")
        audio_measurement = measure_final_loudness(final_temp)
        report = {
            "version": 2, "output": str(output), "sources": [str(p) for p in clips],
            "narration": str(narration) if narration else None,
            "music": str(music) if music else None, "subtitles": str(srt) if srt else None,
            "video": {"width": video["width"], "height": video["height"], "frames": int(video["nb_read_frames"]),
                      "fps": video["avg_frame_rate"], "duration_seconds": duration(video), "codec": video["codec_name"]},
            "audio": {"sample_rate": int(audio["sample_rate"]), "channels": audio["channels"], "codec": audio["codec_name"],
                      "loudness_target": {"integrated_lufs": -16, "true_peak_dbtp": -2, "lra_lu": 7},
                      "measured": audio_measurement, "narration_ducking": narration is not None,
                      "music_gain": 0.12 if music else None, "music_fade_seconds": 0.5 if music else None},
            "subtitle_cues": subtitle_cues,
            "checks": {"six_inputs_at_least_5s_with_audio": True, "each_shot_120_frames": True,
                       "final_720_frames": True, "final_30_seconds": True, "final_48khz_stereo": True,
                       "subtitle_track_verified": bool(srt), "loudness_measured": True},
        }
        # The output path is never an input and pre-existing files are refused.
        if output.exists():
            raise ValueError(f"Output appeared during rendering: {output}; choose a different name.")
        with report_path.open("x", encoding="utf-8") as report_file:
            json.dump(report, report_file, ensure_ascii=False, indent=2)
            report_file.write("\n")
        final_temp.rename(output)
    print(f"Saved: {output}\nReport: {report_path}\nVerified: 1344x768, 24fps, 720 frames, 30.000s; stereo AAC 48kHz.")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (ValueError, RuntimeError, OSError, json.JSONDecodeError) as error:
        print(f"ERROR: {error}", file=sys.stderr)
        raise SystemExit(1)