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

Stage 11: Assembly
- Reads all scene videos from 10-videos/scene_XX/current/video.mp4
- Normalizes clips to consistent format
- Adds timed captions if captions_overlay exists in script
- Concatenates all clips into final video
- Saves to 11-assembly/{iteration}/final.mp4

Usage:
    ./venv/bin/python scripts/assemble_video_v3_kling.py "story-name" --take 003
    ./venv/bin/python scripts/assemble_video_v3_kling.py "story-name" --take 003 --no-captions
"""

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

# Add lib to path
sys.path.insert(0, str(Path(__file__).parent))
from lib import (
    get_story_path,
    get_take_path,
    ensure_dir,
    save_metadata,
    update_current_symlink,
    log_assembly
)

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
    escaped_text = (
        text.replace("\\", "\\\\")
        .replace("'", "\u2019")  # Apostrophe → right single quotation mark
        .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 caption error: {result.stderr[:200]}")
            return False
        return True
    except Exception as e:
        print(f"  Caption error: {e}")
        return False


def normalize_clip(input_path: Path, output_path: Path) -> bool:
    """Normalize clip to consistent format for concatenation."""
    cmd = [
        "ffmpeg", "-y",
        "-i", str(input_path),
        "-c:v", "libx264", "-preset", "fast", "-crf", "18",
        "-c:a", "aac", "-b:a", "192k",
        "-pix_fmt", "yuv420p",  # Ensure consistent pixel format
        str(output_path)
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
        if result.returncode != 0:
            print(f"  FFmpeg normalize error: {result.stderr[:200]}")
            return False
        return True
    except Exception as e:
        print(f"  Normalize 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.txt"
    with open(concat_file, 'w') as f:
        for path in video_paths:
            f.write(f"file '{path.resolve()}'\n")

    cmd = [
        "ffmpeg", "-y",
        "-f", "concat", "-safe", "0",
        "-i", str(concat_file),
        "-c", "copy",  # Stream copy since all clips normalized
        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[:500]}")
            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 using ffprobe."""
    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:
        try:
            return float(result.stdout.strip())
        except ValueError:
            return 0.0
    return 0.0


def get_video_resolution(path: Path) -> tuple:
    """Get video resolution (width, height) using ffprobe."""
    cmd = [
        "ffprobe", "-v", "error",
        "-select_streams", "v:0",
        "-show_entries", "stream=width,height",
        "-of", "csv=s=x:p=0",
        str(path)
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode == 0:
        try:
            w, h = result.stdout.strip().split('x')
            return (int(w), int(h))
        except:
            return (0, 0)
    return (0, 0)


def assemble_video(
    story_name: str,
    take: str = "current",
    no_captions: bool = False
) -> bool:
    """
    Assemble final video from scene clips.

    Args:
        story_name: Story name
        take: Take number or "current"
        no_captions: Skip caption overlay

    Returns:
        True if successful
    """
    # Get paths
    try:
        story_path = get_story_path(story_name)
        take_path = get_take_path(story_name, take)
    except Exception as e:
        print(f"Error: {e}")
        return False

    # Get take number
    if take == "current":
        current_link = story_path / "current"
        if not current_link.is_symlink():
            print(f"Error: No current take set")
            return False
        take_num = current_link.readlink().name
    else:
        take_num = take

    print(f"Assembling video:")
    print(f"  Story: {story_name}")
    print(f"  Take: {take_num}")
    print()

    # Read adapted script for caption info
    script_path = take_path / "08-script-adapted" / "current" / "script.json"
    if not script_path.exists():
        print(f"Error: Adapted script not found: {script_path}")
        return False

    with open(script_path) as f:
        script = json.load(f)

    scenes = script.get("scenes", [])
    if not scenes:
        print("Error: No scenes in script")
        return False

    # Create assembly output directory
    assembly_stage_path = take_path / "11-assembly"
    ensure_dir(assembly_stage_path)

    # Get next iteration number
    existing_iterations = []
    for item in assembly_stage_path.iterdir():
        if item.is_dir() and item.name.isdigit() and len(item.name) == 3:
            existing_iterations.append(int(item.name))

    iteration_num = max(existing_iterations) + 1 if existing_iterations else 1
    iteration = f"{iteration_num:03d}"
    iter_path = assembly_stage_path / iteration
    ensure_dir(iter_path)

    print(f"Assembly iteration: {iteration}")
    print()

    # Build caption timing map from script (if captions_overlay exists)
    # Each scene can have multiple caption lines
    caption_timing = {}  # {scene_num: [caption1, caption2, ...]}
    if not no_captions and "captions_overlay" in script:
        for item in script["captions_overlay"]:
            scene_num = item["scene"]
            text = item["text"]
            timing = item["timing"]  # e.g., "0-4s"

            try:
                start, end = timing.replace("s", "").split("-")
                caption = {
                    "text": text,
                    "start": float(start),
                    "end": float(end),
                    "position": item.get("position", "bottom")
                }

                # Append to list of captions for this scene
                if scene_num not in caption_timing:
                    caption_timing[scene_num] = []
                caption_timing[scene_num].append(caption)

            except Exception as e:
                print(f"Warning: Invalid caption timing for scene {scene_num}: {timing}")

    # Track cumulative time for logging
    cumulative_time = 0.0

    # Process each scene
    processed_clips = []
    temp_dir = iter_path / "temp"
    ensure_dir(temp_dir)

    for scene in scenes:
        scene_num = scene["number"]

        # Get scene video
        scene_video_dir = take_path / "10-videos" / f"scene_{scene_num:02d}" / "current"
        if not scene_video_dir.is_symlink():
            print(f"Error: No current video for scene {scene_num}")
            return False

        video_path = scene_video_dir / "video.mp4"
        if not video_path.exists():
            print(f"Error: Video not found for scene {scene_num}: {video_path}")
            return False

        # Get video duration
        duration = get_video_duration(video_path)

        # Process clip
        if scene_num in caption_timing:
            # Add captions (can be multiple lines)
            captions = caption_timing[scene_num]
            print(f"Scene {scene_num}: {len(captions)} caption line(s) [{duration:.1f}s]")

            # Apply captions sequentially
            current_input = video_path
            success = True

            for idx, ct in enumerate(captions):
                output_captioned = temp_dir / f"scene_{scene_num:02d}_captioned_{idx + 1}.mp4"

                print(f"  Line {idx + 1}/{len(captions)}: '{ct['text']}' ({ct['start']}-{ct['end']}s)")

                if add_timed_caption(
                    current_input, output_captioned,
                    ct["text"], ct["start"], ct["end"],
                    position=ct.get("position", "bottom")
                ):
                    current_input = output_captioned
                else:
                    print(f"  Caption failed for line {idx + 1}")
                    success = False
                    break

            if success:
                processed_clips.append(current_input)
            else:
                print(f"  Caption processing failed, normalizing original")
                output_normalized = temp_dir / f"scene_{scene_num:02d}_normalized.mp4"
                if normalize_clip(video_path, output_normalized):
                    processed_clips.append(output_normalized)
                else:
                    print(f"  Normalization failed, using original (may cause issues)")
                    processed_clips.append(video_path)
        else:
            # No caption, just normalize
            output_normalized = temp_dir / f"scene_{scene_num:02d}_normalized.mp4"
            print(f"Scene {scene_num}: No caption [{duration:.1f}s]")

            if normalize_clip(video_path, output_normalized):
                processed_clips.append(output_normalized)
            else:
                print(f"  Normalization failed, using original (may cause issues)")
                processed_clips.append(video_path)

        cumulative_time += duration

    if not processed_clips:
        print("Error: No clips to assemble")
        return False

    # Concatenate all clips
    output_path = iter_path / "final.mp4"
    print(f"\nAssembling {len(processed_clips)} clips...")

    if not concatenate_videos(processed_clips, output_path):
        print("Error: Failed to concatenate videos")
        return False

    # Get final video stats
    final_duration = get_video_duration(output_path)
    final_resolution = get_video_resolution(output_path)
    size_mb = output_path.stat().st_size / 1024 / 1024

    print()
    print("=" * 60)
    print("Final video assembled!")
    print("=" * 60)
    print(f"Output: {output_path}")
    print(f"Duration: {final_duration:.1f}s")
    print(f"Resolution: {final_resolution[0]}x{final_resolution[1]}")
    print(f"Size: {size_mb:.1f}MB")
    print(f"Scenes: {len(scenes)}")
    if caption_timing:
        print(f"Captions: {len(caption_timing)}")
    print()

    # Save metadata
    metadata = {
        "stage": "11-assembly",
        "take": take_num,
        "iteration": iteration,
        "timestamp": datetime.utcnow().isoformat() + "Z",
        "source_script": "../08-script-adapted/current/script.json",
        "scene_count": len(scenes),
        "scenes_used": [s["number"] for s in scenes],
        "caption_count": len(caption_timing),
        "captioned_scenes": list(caption_timing.keys()) if caption_timing else [],
        "final_duration": final_duration,
        "final_resolution": f"{final_resolution[0]}x{final_resolution[1]}",
        "final_size_mb": round(size_mb, 2),
        "output_file": "final.mp4",
        "concat_method": "ffmpeg_concat_demuxer",
        "normalization_applied": True
    }

    metadata_file = iter_path / "metadata.json"
    save_metadata(metadata_file, metadata)
    print(f"✓ Saved metadata: {metadata_file}")

    # Update current symlink
    update_current_symlink(assembly_stage_path, iteration)
    print(f"✓ Updated current → {iteration}")

    # Log to pipeline.log
    log_assembly(
        story_name,
        take_num,
        iteration,
        len(scenes),
        final_duration,
        size_mb
    )

    return True


def main():
    parser = argparse.ArgumentParser(
        description="Assemble final video from scene clips with optional captions"
    )
    parser.add_argument("story_name", help="Story name")
    parser.add_argument("--take", default="current", help="Take number (default: current)")
    parser.add_argument("--no-captions", action="store_true", help="Skip caption overlay")

    args = parser.parse_args()

    success = assemble_video(
        args.story_name,
        args.take,
        args.no_captions
    )

    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()
