#!/usr/bin/env python3
"""
Assemble scene clips into final video using FFmpeg.

Supports simple concat or crossfade transitions.
"""

import os
import sys
import json
import subprocess
from pathlib import Path


def get_video_duration(video_path: Path) -> float:
    """Get video duration in seconds."""
    result = subprocess.run([
        "ffprobe", "-v", "error",
        "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1",
        str(video_path)
    ], capture_output=True, text=True)
    return float(result.stdout.strip())


def assemble_simple_concat(clips: list, output_path: Path) -> bool:
    """Simple concatenation without transitions."""
    # Create concat file
    concat_file = output_path.parent / "concat_list.txt"
    with open(concat_file, "w") as f:
        for clip in clips:
            f.write(f"file '{clip}'\n")

    cmd = [
        "ffmpeg", "-y",
        "-f", "concat",
        "-safe", "0",
        "-i", str(concat_file),
        "-c", "copy",
        str(output_path)
    ]

    print(f"Running: {' '.join(cmd)}")
    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        print(f"Error: {result.stderr}")
        return False

    return True


def assemble_with_crossfades(clips: list, output_path: Path, fade_duration: float = 0.5) -> bool:
    """Concatenation with crossfade transitions between clips. Includes audio crossfades."""
    if len(clips) < 2:
        return assemble_simple_concat(clips, output_path)

    # Build complex filter for crossfades
    # Each clip needs to be trimmed and faded
    n = len(clips)

    # Input files
    inputs = []
    for clip in clips:
        inputs.extend(["-i", str(clip)])

    # Build filter_complex
    filter_parts = []

    # First, set up all inputs with proper formatting
    for i in range(n):
        filter_parts.append(f"[{i}:v]setpts=PTS-STARTPTS[v{i}]")
        filter_parts.append(f"[{i}:a]asetpts=PTS-STARTPTS[a{i}]")

    # Now create crossfades between consecutive clips (video and audio)
    # xfade filter: [v0][v1]xfade=transition=fade:duration=0.5:offset=7.5[vout01]
    # acrossfade: [a0][a1]acrossfade=d=0.5:c1=tri:c2=tri[aout01]
    prev_v_output = "v0"
    prev_a_output = "a0"

    for i in range(1, n):
        # Calculate offset (when the transition starts)
        # offset = sum of previous clip durations minus fade overlaps
        offset = sum(get_video_duration(Path(clips[j])) for j in range(i)) - (i * fade_duration)

        if i == n - 1:
            # Last transition outputs to final
            v_output_label = "vout"
            a_output_label = "aout"
        else:
            v_output_label = f"vout{i}"
            a_output_label = f"aout{i}"

        # Video crossfade
        filter_parts.append(
            f"[{prev_v_output}][v{i}]xfade=transition=fade:duration={fade_duration}:offset={offset:.2f}[{v_output_label}]"
        )

        # Audio crossfade
        filter_parts.append(
            f"[{prev_a_output}][a{i}]acrossfade=d={fade_duration}:c1=tri:c2=tri[{a_output_label}]"
        )

        prev_v_output = v_output_label
        prev_a_output = a_output_label

    filter_complex = ";".join(filter_parts)

    cmd = [
        "ffmpeg", "-y",
        *inputs,
        "-filter_complex", filter_complex,
        "-map", "[vout]",
        "-map", "[aout]",
        "-c:v", "libx264",
        "-preset", "medium",
        "-crf", "18",
        "-pix_fmt", "yuv420p",
        "-c:a", "aac",
        "-b:a", "192k",
        str(output_path)
    ]

    print(f"Running FFmpeg with crossfades (video + audio)...")
    print(f"Filter: {filter_complex[:200]}...")

    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        print(f"Error: {result.stderr[:500]}")
        return False

    return True


def assemble_with_fade_black(clips: list, output_path: Path, fade_duration: float = 0.3) -> bool:
    """Each clip fades to black, then next clip fades from black. Includes audio with fades."""
    if len(clips) < 1:
        return False

    n = len(clips)

    # Input files
    inputs = []
    for clip in clips:
        inputs.extend(["-i", str(clip)])

    # Build filter: fade out each clip (video AND audio), then concat
    filter_parts = []

    for i in range(n):
        duration = get_video_duration(Path(clips[i]))
        fade_start = duration - fade_duration

        # Video: fade in at start, fade out at end
        filter_parts.append(
            f"[{i}:v]fade=t=in:st=0:d={fade_duration},fade=t=out:st={fade_start:.2f}:d={fade_duration}[v{i}]"
        )

        # Audio: fade in at start, fade out at end
        filter_parts.append(
            f"[{i}:a]afade=t=in:st=0:d={fade_duration},afade=t=out:st={fade_start:.2f}:d={fade_duration}[a{i}]"
        )

    # Concat all processed clips (video and audio)
    # concat expects interleaved pairs: [v0][a0][v1][a1]...
    concat_inputs = "".join(f"[v{i}][a{i}]" for i in range(n))
    filter_parts.append(f"{concat_inputs}concat=n={n}:v=1:a=1[vout][aout]")

    filter_complex = ";".join(filter_parts)

    cmd = [
        "ffmpeg", "-y",
        *inputs,
        "-filter_complex", filter_complex,
        "-map", "[vout]",
        "-map", "[aout]",
        "-c:v", "libx264",
        "-preset", "medium",
        "-crf", "18",
        "-pix_fmt", "yuv420p",
        "-c:a", "aac",
        "-b:a", "192k",
        str(output_path)
    ]

    print(f"Running FFmpeg with fade transitions (video + audio)...")

    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        print(f"Error: {result.stderr[:500]}")
        return False

    return True


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Assemble scene clips into final video")
    parser.add_argument("task_dir", help="Task directory path")
    parser.add_argument("--mode", "-m", choices=["concat", "crossfade", "fade"],
                       default="fade", help="Assembly mode (default: fade)")
    parser.add_argument("--fade-duration", "-f", type=float, default=0.3,
                       help="Fade duration in seconds (default: 0.3)")
    parser.add_argument("--output", "-o", help="Output filename (default: final_video.mp4)")
    args = parser.parse_args()

    task_dir = Path(args.task_dir)
    scenes_dir = task_dir / "scenes"

    # Find all clips
    clips = []
    for scene_dir in sorted(scenes_dir.glob("scene_*")):
        clip_path = scene_dir / "clip.mp4"
        if clip_path.exists():
            clips.append(str(clip_path))

    if not clips:
        print("Error: No clips found")
        sys.exit(1)

    print(f"Found {len(clips)} clips")

    # Output path
    assembly_dir = task_dir / "assembly"
    assembly_dir.mkdir(exist_ok=True)

    output_name = args.output or "final_video.mp4"
    output_path = assembly_dir / output_name

    # Assemble
    print(f"\nAssembling with mode: {args.mode}")
    print(f"Output: {output_path}\n")

    if args.mode == "concat":
        success = assemble_simple_concat(clips, output_path)
    elif args.mode == "crossfade":
        success = assemble_with_crossfades(clips, output_path, args.fade_duration)
    else:  # fade
        success = assemble_with_fade_black(clips, output_path, args.fade_duration)

    if success:
        size_mb = output_path.stat().st_size / (1024 * 1024)
        duration = get_video_duration(output_path)
        print(f"\nSuccess!")
        print(f"Output: {output_path}")
        print(f"Size: {size_mb:.1f} MB")
        print(f"Duration: {duration:.1f}s")
    else:
        print("\nAssembly failed")
        sys.exit(1)


if __name__ == "__main__":
    main()
