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

Reads video_prompts.json for:
- Clip order and paths
- Caption text and timing from captions_overlay

Uses FFmpeg to:
1. Add timed captions to clips that have them
2. Concatenate all clips into final video
"""

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

sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))

from dotenv import load_dotenv
load_dotenv(Path(__file__).parent.parent.parent / ".env")


def add_timed_caption(
    input_path: Path,
    output_path: Path,
    text: str,
    start_time: float,
    end_time: float,
    font_size: int = 42,
    font_color: str = "white",
    bg_color: str = "black@0.7",
    position: str = "bottom",
) -> bool:
    """Add timed caption overlay to video using FFmpeg."""

    # Escape special characters for FFmpeg drawtext filter
    # FFmpeg needs: ' -> '\'' for shell, and special chars escaped with backslash
    escaped_text = text.replace("\\", "\\\\").replace("'", "\u2019").replace(":", "\\:").replace("%", "\\%")

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

    # FFmpeg drawtext filter with enable for timing
    filter_str = (
        f"drawtext=text='{escaped_text}':"
        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=15:"
        f"fontfile=/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf:"
        f"enable='between(t,{start_time},{end_time})'"
    )

    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, reencode: bool = True) -> bool:
    """Concatenate videos using FFmpeg."""

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

    if reencode:
        # Re-encode for consistent format
        cmd = [
            "ffmpeg", "-y",
            "-f", "concat", "-safe", "0",
            "-i", str(concat_file),
            "-c:v", "libx264", "-preset", "fast", "-crf", "18",
            "-c:a", "aac", "-b:a", "192k",
            str(output_path)
        ]
    else:
        # Stream copy (faster but may have issues with different formats)
        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=600)
        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 get_video_duration(path: Path) -> float:
    """Get video duration in seconds."""
    cmd = [
        "ffprobe", "-v", "error",
        "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1",
        str(path)
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode == 0:
        return float(result.stdout.strip())
    return 0.0


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Assemble video with timed captions")
    parser.add_argument("prompts_file", help="Path to video_prompts.json")
    parser.add_argument("--videos-dir", "-v", help="Videos directory (default: from prompts)")
    parser.add_argument("--output", "-o", help="Output file path")
    parser.add_argument("--no-captions", action="store_true", help="Skip caption overlay")
    parser.add_argument("--no-reencode", action="store_true", help="Use stream copy instead of re-encoding")
    args = parser.parse_args()

    # Load config
    prompts_path = Path(args.prompts_file)
    if not prompts_path.exists():
        print(f"Error: Prompts file not found: {prompts_path}")
        sys.exit(1)

    with open(prompts_path) as f:
        config = json.load(f)

    # Videos directory
    if args.videos_dir:
        videos_dir = Path(args.videos_dir)
    else:
        # Infer from prompts location
        videos_dir = prompts_path.parent.parent.parent / "videos" / "002"

    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 = prompts_path.parent.parent.parent / "assembly" / "001" / "final.mp4"

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

    print(f"Videos: {videos_dir}")
    print(f"Output: {output_path}")

    # Build caption timing map
    caption_timing = {}
    for item in config.get("captions_overlay", []):
        clip_num = item["clip"]
        text = item["text"]
        timing = item["timing"]  # e.g., "0-4s"
        start, end = timing.replace("s", "").split("-")
        caption_timing[clip_num] = {
            "text": text,
            "start": float(start),
            "end": float(end)
        }

    # Process each clip
    processed_clips = []

    for clip in config["clips"]:
        clip_id = clip["clip_id"]
        input_path = videos_dir / f"clip_{clip_id:02d}" / "video.mp4"

        if not input_path.exists():
            print(f"Warning: Clip {clip_id} not found: {input_path}")
            continue

        clip_dir = videos_dir / f"clip_{clip_id:02d}"

        if not args.no_captions and clip_id in caption_timing:
            ct = caption_timing[clip_id]
            output_captioned = clip_dir / "video_captioned.mp4"

            print(f"Clip {clip_id}: '{ct['text'][:35]}...' ({ct['start']}-{ct['end']}s)")

            if add_timed_caption(
                input_path, output_captioned,
                ct["text"], ct["start"], ct["end"]
            ):
                processed_clips.append(output_captioned)
            else:
                print(f"  Caption failed, using original")
                processed_clips.append(input_path)
        else:
            # Re-encode non-captioned clips for consistent format
            output_normalized = clip_dir / "video_normalized.mp4"
            print(f"Clip {clip_id}: Normalizing (no caption)")

            cmd = [
                "ffmpeg", "-y",
                "-i", str(input_path),
                "-c:v", "libx264", "-preset", "fast", "-crf", "18",
                "-c:a", "aac", "-b:a", "192k",
                str(output_normalized)
            ]
            try:
                result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
                if result.returncode == 0:
                    processed_clips.append(output_normalized)
                else:
                    print(f"  Normalize failed, using original")
                    processed_clips.append(input_path)
            except Exception as e:
                print(f"  Error: {e}")
                processed_clips.append(input_path)

    if not processed_clips:
        print("Error: No clips to assemble")
        sys.exit(1)

    # Concatenate
    print(f"\nAssembling {len(processed_clips)} clips...")
    if concatenate_videos(processed_clips, output_path, reencode=not args.no_reencode):
        duration = get_video_duration(output_path)
        size_mb = output_path.stat().st_size / 1024 / 1024
        print(f"\n✓ Final video: {output_path}")
        print(f"  Duration: {duration:.1f}s")
        print(f"  Size: {size_mb:.1f}MB")
    else:
        print("Error: Failed to assemble video")
        sys.exit(1)


if __name__ == "__main__":
    main()
