backend/services/video_assembler.py

"""Video assembly service - combines all elements into final export."""
import os
import subprocess
import logging
import json
from pathlib import Path
from typing import List, Optional, Dict
from dataclasses import dataclass

logger = logging.getLogger(__name__)


@dataclass
class AssemblyConfig:
    """Configuration for video assembly."""
    clip_start: float
    clip_end: float
    source_video: str
    output_path: str
    voiceover_path: Optional[str] = None
    subtitle_path: Optional[str] = None
    stock_images: List[str] = None
    sound_effects: List[Dict] = None
    music_path: Optional[str] = None
    color_grade: str = "teal_orange"
    include_original_audio: bool = False
    intro_card_path: Optional[str] = None
    title: str = ""
    
    def __post_init__(self):
        if self.stock_images is None:
            self.stock_images = []
        if self.sound_effects is None:
            self.sound_effects = []


# Color grading LUT filter definitions (ffmpeg eq/colorbalance approximations)
COLOR_GRADES = {
    "teal_orange": "colorbalance=rs=0.1:gs=-0.1:bs=-0.1:rm=0.1:gm=0:bm=-0.1:rh=0.05:gh=-0.05:bh=-0.1,eq=saturation=1.3",
    "high_contrast_bw": "colorchannelmixer=.3:.4:.3:0:.3:.4:.3:0:.3:.4:.3,eq=contrast=1.4:brightness=0.05",
    "warm_golden": "colorbalance=rs=0.15:gs=0.05:bs=-0.1:rm=0.1:gm=0.05:bm=-0.05,eq=saturation=1.2:brightness=0.05",
    "cold_blue": "colorbalance=rs=-0.1:gs=-0.05:bs=0.15:rm=-0.05:gm=0:bm=0.1,eq=saturation=1.1:contrast=1.1",
    "vivid_oversaturated": "eq=saturation=1.8:contrast=1.2:brightness=0.03,unsharp=5:5:1.0",
}


def trim_clip(source_video: str, start: float, end: float, output_path: str) -> str:
    """Trim a clip from the source video.
    
    Args:
        source_video: Path to source video
        start: Start time in seconds
        end: End time in seconds
        output_path: Output file path
    
    Returns:
        Path to trimmed clip
    """
    cmd = [
        "ffmpeg", "-y",
        "-ss", str(start),
        "-i", source_video,
        "-t", str(end - start),
        "-c:v", "libx264", "-crf", "18",
        "-c:a", "aac", "-b:a", "192k",
        "-preset", "fast",
        output_path
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
    if result.returncode != 0:
        raise RuntimeError(f"Trim failed: {result.stderr[:200]}")
    
    return output_path


def resize_to_portrait(input_path: str, output_path: str, width: int = 1080, height: int = 1920) -> str:
    """Resize and crop video to portrait (9:16) format.
    
    Center-crops the video to fit 1080x1920.
    """
    # Scale to fill and center crop
    filter_str = (
        f"scale={width}:{height}:force_original_aspect_ratio=increase,"
        f"crop={width}:{height}"
    )
    
    cmd = [
        "ffmpeg", "-y",
        "-i", input_path,
        "-vf", filter_str,
        "-c:v", "libx264", "-crf", "18",
        "-c:a", "copy",
        "-preset", "fast",
        output_path
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
    if result.returncode != 0:
        raise RuntimeError(f"Resize failed: {result.stderr[:200]}")
    
    return output_path


def apply_color_grade(input_path: str, output_path: str, grade_name: str) -> str:
    """Apply color grading to video.
    
    Args:
        input_path: Input video path
        output_path: Output path
        grade_name: Color grade preset name
    
    Returns:
        Path to graded video
    """
    grade_filter = COLOR_GRADES.get(grade_name, COLOR_GRADES["teal_orange"])
    
    cmd = [
        "ffmpeg", "-y",
        "-i", input_path,
        "-vf", grade_filter,
        "-c:v", "libx264", "-crf", "18",
        "-c:a", "copy",
        "-preset", "fast",
        output_path
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
    if result.returncode != 0:
        logger.warning(f"Color grade failed, using original: {result.stderr[:200]}")
        return input_path
    
    return output_path


def apply_vignette_and_grain(input_path: str, output_path: str) -> str:
    """Apply vignette and film grain effects."""
    filter_str = (
        "vignette=PI/5,"
        "noise=c0s=8:c0f=t+u"  # Light film grain
    )
    
    cmd = [
        "ffmpeg", "-y",
        "-i", input_path,
        "-vf", filter_str,
        "-c:v", "libx264", "-crf", "18",
        "-c:a", "copy",
        "-preset", "fast",
        output_path
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
    if result.returncode != 0:
        logger.warning(f"Effects failed: {result.stderr[:200]}")
        return input_path
    
    return output_path


def create_ken_burns_clip(image_path: str, duration: float, output_path: str) -> str:
    """Create a Ken Burns effect video from a still image.
    
    Applies slow zoom and pan over the image.
    """
    fps = 30
    total_frames = int(duration * fps)
    
    # Zoompan filter: slow zoom in from 1.0x to 1.3x over duration
    filter_str = (
        f"zoompan=z='min(zoom+0.001,1.3)':x='iw/2-(iw/zoom/2)':y='ih/2-(ih/zoom/2)'"
        f":d={total_frames}:s=1080x1920:fps={fps}"
    )
    
    cmd = [
        "ffmpeg", "-y",
        "-loop", "1",
        "-i", image_path,
        "-vf", filter_str,
        "-t", str(duration),
        "-c:v", "libx264", "-crf", "20",
        "-pix_fmt", "yuv420p",
        "-preset", "fast",
        output_path
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
    if result.returncode != 0:
        logger.warning(f"Ken Burns failed: {result.stderr[:200]}")
        # Create simple static image video as fallback
        cmd_fallback = [
            "ffmpeg", "-y", "-loop", "1", "-i", image_path,
            "-vf", f"scale=1080:1920:force_original_aspect_ratio=decrease,pad=1080:1920:(ow-iw)/2:(oh-ih)/2",
            "-t", str(duration), "-c:v", "libx264", "-pix_fmt", "yuv420p", output_path
        ]
        subprocess.run(cmd_fallback, capture_output=True, timeout=30)
    
    return output_path


def create_intro_card(
    title: str,
    keyframe_path: Optional[str],
    output_path: str,
    duration: float = 3.0,
    style: str = "heatwave"
) -> str:
    """Create an animated intro card.
    
    Args:
        title: Clip title text
        keyframe_path: Background keyframe image (will be blurred)
        output_path: Output video path
        duration: Intro duration in seconds
        style: Visual style matching caption style
    
    Returns:
        Path to intro card video
    """
    from PIL import Image, ImageDraw, ImageFont, ImageFilter
    import tempfile
    
    width, height = 1080, 1920
    fps = 30
    frames_dir = Path(output_path).parent / "intro_frames"
    frames_dir.mkdir(exist_ok=True)
    
    # Create background
    if keyframe_path and os.path.exists(keyframe_path):
        bg = Image.open(keyframe_path).resize((width, height)).filter(ImageFilter.GaussianBlur(20))
    else:
        # Gradient background
        bg = Image.new("RGB", (width, height), (20, 20, 40))
        draw = ImageDraw.Draw(bg)
        for y in range(height):
            ratio = y / height
            r = int(20 + 30 * ratio)
            g = int(20 + 10 * ratio)
            b = int(40 + 60 * ratio)
            draw.line([(0, y), (width, y)], fill=(r, g, b))
    
    # Darken background
    overlay = Image.new("RGBA", (width, height), (0, 0, 0, 150))
    bg = bg.convert("RGBA")
    bg = Image.alpha_composite(bg, overlay).convert("RGB")
    
    # Load font
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 72)
        font_small = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 36)
    except:
        font = ImageFont.load_default()
        font_small = font
    
    # Generate frames with text animation (slide up + fade in)
    total_frames = int(duration * fps)
    
    for frame_idx in range(total_frames):
        frame = bg.copy()
        draw = ImageDraw.Draw(frame)
        
        # Animation progress (0 to 1)
        progress = min(frame_idx / (fps * 1.0), 1.0)  # Animate over first 1 second
        
        # Text slide up effect
        y_offset = int(50 * (1 - progress))  # Slide up 50px
        alpha = int(255 * progress)
        
        # Draw title (word wrap)
        words = title.split()
        lines = []
        current_line = ""
        for word in words:
            test_line = f"{current_line} {word}".strip()
            bbox = draw.textbbox((0, 0), test_line, font=font)
            if bbox[2] - bbox[0] > width - 100:
                lines.append(current_line)
                current_line = word
            else:
                current_line = test_line
        if current_line:
            lines.append(current_line)
        
        # Center text vertically
        line_height = 90
        total_text_height = len(lines) * line_height
        start_y = (height - total_text_height) // 2 + y_offset
        
        for i, line in enumerate(lines):
            bbox = draw.textbbox((0, 0), line, font=font)
            text_width = bbox[2] - bbox[0]
            x = (width - text_width) // 2
            y = start_y + i * line_height
            
            # Draw text shadow
            draw.text((x + 3, y + 3), line, font=font, fill=(0, 0, 0))
            # Draw text
            draw.text((x, y), line, font=font, fill=(255, 255, 255))
        
        # Save frame
        frame_path = str(frames_dir / f"frame_{frame_idx:05d}.png")
        frame.save(frame_path)
    
    # Assemble frames into video
    cmd = [
        "ffmpeg", "-y",
        "-framerate", str(fps),
        "-i", str(frames_dir / "frame_%05d.png"),
        "-c:v", "libx264", "-crf", "18",
        "-pix_fmt", "yuv420p",
        "-preset", "fast",
        output_path
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
    
    # Clean up frames
    import shutil
    shutil.rmtree(frames_dir, ignore_errors=True)
    
    if result.returncode != 0:
        logger.error(f"Intro card creation failed: {result.stderr[:200]}")
        # Create a simple colored clip as fallback
        cmd_fallback = [
            "ffmpeg", "-y", "-f", "lavfi",
            "-i", f"color=c=black:s={width}x{height}:d={duration}:r={fps}",
            "-vf", f"drawtext=text='{title[:30]}':fontsize=60:fontcolor=white:x=(w-text_w)/2:y=(h-text_h)/2",
            "-c:v", "libx264", "-pix_fmt", "yuv420p",
            output_path
        ]
        subprocess.run(cmd_fallback, capture_output=True, timeout=30)
    
    return output_path


def mix_audio(
    voiceover_path: str,
    output_path: str,
    original_audio_path: Optional[str] = None,
    music_path: Optional[str] = None,
    sfx_paths: List[Dict] = None,
    clip_duration: float = 60.0
) -> str:
    """Mix all audio tracks together.
    
    Audio levels:
    - Voiceover: 0dB (primary)
    - Original audio: -18dB
    - Background music: -20dB (ducked under voiceover)
    - Sound effects: -12dB
    
    Args:
        voiceover_path: Path to voiceover audio
        output_path: Output mixed audio path
        original_audio_path: Optional original video audio
        music_path: Optional background music
        sfx_paths: List of {path, offset_seconds} for sound effects
        clip_duration: Target clip duration
    
    Returns:
        Path to mixed audio file
    """
    if sfx_paths is None:
        sfx_paths = []
    
    inputs = ["-i", voiceover_path]
    filter_parts = []
    input_idx = 0
    
    # Voiceover is always input [0]
    filter_parts.append(f"[0:a]volume=1.0,apad=pad_dur=0[vo]")
    input_idx = 1
    
    mix_inputs = ["[vo]"]
    
    # Add original audio if provided
    if original_audio_path and os.path.exists(original_audio_path):
        inputs.extend(["-i", original_audio_path])
        filter_parts.append(f"[{input_idx}:a]volume=0.125[orig]")  # -18dB
        mix_inputs.append("[orig]")
        input_idx += 1
    
    # Add background music if provided
    if music_path and os.path.exists(music_path):
        inputs.extend(["-i", music_path])
        filter_parts.append(f"[{input_idx}:a]volume=0.1,aloop=loop=-1:size=2e+09,atrim=0:{clip_duration}[music]")  # -20dB, looped
        mix_inputs.append("[music]")
        input_idx += 1
    
    # Add sound effects
    for i, sfx in enumerate(sfx_paths[:5]):  # Max 5 SFX
        sfx_file = sfx.get("path", "")
        offset = sfx.get("offset_seconds", 0)
        if os.path.exists(sfx_file):
            inputs.extend(["-i", sfx_file])
            filter_parts.append(f"[{input_idx}:a]volume=0.25,adelay={int(offset*1000)}|{int(offset*1000)}[sfx{i}]")  # -12dB
            mix_inputs.append(f"[sfx{i}]")
            input_idx += 1
    
    # Mix all inputs
    num_inputs = len(mix_inputs)
    mix_str = "".join(mix_inputs)
    filter_parts.append(f"{mix_str}amix=inputs={num_inputs}:duration=first:dropout_transition=2[out]")
    
    filter_complex = ";".join(filter_parts)
    
    cmd = [
        "ffmpeg", "-y",
        *inputs,
        "-filter_complex", filter_complex,
        "-map", "[out]",
        "-c:a", "aac", "-b:a", "192k",
        "-t", str(clip_duration),
        output_path
    ]
    
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
        if result.returncode == 0:
            return output_path
        else:
            logger.warning(f"Audio mix failed: {result.stderr[:200]}")
            # Fallback: just use voiceover
            subprocess.run(
                ["ffmpeg", "-y", "-i", voiceover_path, "-c:a", "aac", "-b:a", "192k", output_path],
                capture_output=True, timeout=30
            )
            return output_path
    except Exception as e:
        logger.error(f"Audio mix error: {e}")
        return voiceover_path


def assemble_final_video(config: AssemblyConfig) -> str:
    """Full video assembly pipeline.
    
    Steps:
    1. Trim source video
    2. Resize to portrait
    3. Apply color grading
    4. Apply effects (vignette + grain)
    5. Burn subtitles
    6. Mix audio
    7. Combine video + audio
    8. Prepend intro card
    9. Export final
    
    Args:
        config: AssemblyConfig with all parameters
    
    Returns:
        Path to final assembled video
    """
    work_dir = Path(config.output_path).parent
    work_dir.mkdir(parents=True, exist_ok=True)
    
    clip_duration = config.clip_end - config.clip_start
    
    # Step 1: Trim
    logger.info(f"Step 1: Trimming {config.clip_start:.1f}s - {config.clip_end:.1f}s")
    trimmed = str(work_dir / "01_trimmed.mp4")
    trim_clip(config.source_video, config.clip_start, config.clip_end, trimmed)
    
    # Step 2: Resize to portrait
    logger.info("Step 2: Resizing to 1080x1920")
    resized = str(work_dir / "02_resized.mp4")
    resize_to_portrait(trimmed, resized)
    
    # Step 3: Color grading
    logger.info(f"Step 3: Applying color grade: {config.color_grade}")
    graded = str(work_dir / "03_graded.mp4")
    apply_color_grade(resized, graded, config.color_grade)
    
    # Step 4: Vignette + grain
    logger.info("Step 4: Applying vignette and grain")
    effected = str(work_dir / "04_effects.mp4")
    apply_vignette_and_grain(graded, effected)
    
    # Step 5: Burn subtitles
    current_video = effected
    if config.subtitle_path and os.path.exists(config.subtitle_path):
        logger.info("Step 5: Burning subtitles")
        subtitled = str(work_dir / "05_subtitled.mp4")
        current_video = burn_subtitles(effected, config.subtitle_path, subtitled)
    
    # Step 6: Mix audio
    logger.info("Step 6: Mixing audio")
    mixed_audio = str(work_dir / "06_mixed_audio.aac")
    
    original_audio = None
    if config.include_original_audio:
        original_audio = str(work_dir / "original_audio.aac")
        subprocess.run(
            ["ffmpeg", "-y", "-i", trimmed, "-vn", "-c:a", "aac", original_audio],
            capture_output=True, timeout=30
        )
    
    if config.voiceover_path and os.path.exists(config.voiceover_path):
        mix_audio(
            voiceover_path=config.voiceover_path,
            output_path=mixed_audio,
            original_audio_path=original_audio,
            music_path=config.music_path,
            sfx_paths=config.sound_effects,
            clip_duration=clip_duration
        )
    else:
        # Just use original audio
        if original_audio:
            mixed_audio = original_audio
        else:
            # Generate silence
            subprocess.run(
                ["ffmpeg", "-y", "-f", "lavfi", "-i", f"anullsrc=r=44100:cl=stereo",
                 "-t", str(clip_duration), "-c:a", "aac", mixed_audio],
                capture_output=True, timeout=10
            )
    
    # Step 7: Combine video + mixed audio
    logger.info("Step 7: Combining video and audio")
    combined = str(work_dir / "07_combined.mp4")
    cmd = [
        "ffmpeg", "-y",
        "-i", current_video,
        "-i", mixed_audio,
        "-map", "0:v:0", "-map", "1:a:0",
        "-c:v", "copy", "-c:a", "aac",
        "-shortest",
        combined
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
    if result.returncode != 0:
        combined = current_video  # Fallback
    
    # Step 8: Prepend intro card
    final_input = combined
    if config.intro_card_path and os.path.exists(config.intro_card_path):
        logger.info("Step 8: Prepending intro card")
        # Add silent audio to intro card
        intro_with_audio = str(work_dir / "intro_audio.mp4")
        cmd = [
            "ffmpeg", "-y",
            "-i", config.intro_card_path,
            "-f", "lavfi", "-i", "anullsrc=r=44100:cl=stereo",
            "-c:v", "copy", "-c:a", "aac", "-shortest",
            intro_with_audio
        ]
        subprocess.run(cmd, capture_output=True, timeout=30)
        
        # Concatenate intro + main clip
        concat_file = str(work_dir / "concat.txt")
        with open(concat_file, "w") as f:
            f.write(f"file '{intro_with_audio}'\nfile '{combined}'\n")
        
        with_intro = str(work_dir / "08_with_intro.mp4")
        cmd = [
            "ffmpeg", "-y", "-f", "concat", "-safe", "0",
            "-i", concat_file,
            "-c:v", "libx264", "-crf", "18",
            "-c:a", "aac", "-b:a", "192k",
            with_intro
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
        if result.returncode == 0:
            final_input = with_intro
    
    # Step 9: Final export with target settings
    logger.info("Step 9: Final export")
    cmd = [
        "ffmpeg", "-y",
        "-i", final_input,
        "-c:v", "libx264", "-crf", "18",
        "-preset", "medium",
        "-r", "60",  # 60fps
        "-s", "1080x1920",
        "-c:a", "aac", "-b:a", "192k",
        "-movflags", "+faststart",
        config.output_path
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
    if result.returncode != 0:
        logger.error(f"Final export failed: {result.stderr[:300]}")
        # Try simpler export
        subprocess.run(
            ["cp", final_input, config.output_path],
            capture_output=True
        )
    
    # Cleanup intermediate files
    for f in work_dir.glob("0*_*.mp4"):
        try:
            os.remove(f)
        except:
            pass
    
    logger.info(f"Assembly complete: {config.output_path}")
    return config.output_path