#!/usr/bin/env python3
"""
Assemble final video from scene clips with captions.

Uses FFmpeg to:
1. Add captions to each scene
2. Concatenate all scenes into final video
"""

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

# Scene captions (Russian)
CAPTIONS = {
    1: "Обычный дождливый день в аэропорту",
    2: "Косарь делает свою работу",
    3: "Як-42 выруливает на взлёт",
    4: "Работник отошёл в сторону",
    5: "Он приветливо машет самолёту",
    6: "Пассажиры в ужасе!",
    7: "«Опять Семёныч вышел траву покосить...»",
    8: "Вид из иллюминатора...",
    9: "— Ты бы хотя бы косу положил!",
}


def add_caption_to_video(
    input_path: Path,
    output_path: Path,
    caption: str,
    font_size: int = 36,
    font_color: str = "white",
    bg_color: str = "black@0.6",
    position: str = "bottom",
) -> bool:
    """Add caption overlay to video using FFmpeg."""

    # Position mapping
    if position == "bottom":
        y_pos = "h-th-40"
    elif position == "top":
        y_pos = "40"
    else:
        y_pos = "(h-th)/2"

    # FFmpeg drawtext filter
    # Using box for background
    filter_str = (
        f"drawtext=text='{caption}':"
        f"fontsize={font_size}:"
        f"fontcolor={font_color}:"
        f"x=(w-tw)/2:"
        f"y={y_pos}:"
        f"box=1:"
        f"boxcolor={bg_color}:"
        f"boxborderw=10:"
        f"fontfile=/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"
    )

    cmd = [
        "ffmpeg", "-y",
        "-i", str(input_path),
        "-vf", filter_str,
        "-c:a", "copy",
        str(output_path)
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
        if result.returncode != 0:
            print(f"  FFmpeg error: {result.stderr[:200]}")
            return False
        return True
    except Exception as e:
        print(f"  Error: {e}")
        return False


def concatenate_videos(video_paths: list, output_path: Path) -> bool:
    """Concatenate videos using FFmpeg concat demuxer."""

    # Create concat file
    concat_file = output_path.parent / "concat_list.txt"
    with open(concat_file, 'w') as f:
        for path in video_paths:
            f.write(f"file '{path}'\n")

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

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        if result.returncode != 0:
            print(f"  FFmpeg concat error: {result.stderr[:200]}")
            return False
        return True
    except Exception as e:
        print(f"  Concat error: {e}")
        return False


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Assemble final video with captions")
    parser.add_argument("videos_dir", help="Directory with scene videos")
    parser.add_argument("--output", "-o", help="Output file path")
    parser.add_argument("--no-captions", action="store_true", help="Skip caption overlay")
    args = parser.parse_args()

    videos_dir = Path(args.videos_dir)
    if not videos_dir.exists():
        print(f"Error: Videos directory not found: {videos_dir}")
        sys.exit(1)

    # Output path
    if args.output:
        output_path = Path(args.output)
    else:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_path = videos_dir.parent / "assembly" / f"final_{timestamp}.mp4"

    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Process each scene
    captioned_videos = []
    temp_dir = output_path.parent / "temp"
    temp_dir.mkdir(exist_ok=True)

    for scene_num in range(1, 10):
        scene_dir = videos_dir / f"scene_{scene_num:02d}"
        video_path = scene_dir / "video.mp4"

        if not video_path.exists():
            print(f"Warning: Video not found: {video_path}")
            continue

        print(f"Processing scene {scene_num}: {CAPTIONS.get(scene_num, '')}")

        if args.no_captions:
            captioned_videos.append(video_path)
        else:
            # Add caption
            captioned_path = temp_dir / f"scene_{scene_num:02d}_captioned.mp4"
            caption = CAPTIONS.get(scene_num, "")

            if add_caption_to_video(video_path, captioned_path, caption):
                captioned_videos.append(captioned_path)
                print(f"  Captioned: {captioned_path.name}")
            else:
                print(f"  Failed to add caption, using original")
                captioned_videos.append(video_path)

    if not captioned_videos:
        print("Error: No videos to assemble")
        sys.exit(1)

    # Concatenate
    print(f"\nAssembling {len(captioned_videos)} videos...")
    if concatenate_videos(captioned_videos, output_path):
        print(f"\nFinal video: {output_path}")
        # Get duration
        probe_cmd = ["ffprobe", "-v", "error", "-show_entries", "format=duration",
                     "-of", "default=noprint_wrappers=1:nokey=1", str(output_path)]
        result = subprocess.run(probe_cmd, capture_output=True, text=True)
        if result.returncode == 0:
            duration = float(result.stdout.strip())
            print(f"Duration: {duration:.1f} seconds")
    else:
        print("Error: Failed to assemble video")
        sys.exit(1)


if __name__ == "__main__":
    main()
