#!/usr/bin/env python3
"""
Video Transitions Example for FMUS-VID.

This example demonstrates the various transition effects available in FMUS-VID.
"""

import os
import sys
import argparse
from pathlib import Path
import logging
import numpy as np
import cv2
import time

# Add parent directory to path to allow running from examples directory
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from fmusvid.operations.transitions import (
    FadeTransition, DissolveTransition, WipeTransition,
    SlideTransition, PixelateTransition, ZoomTransition,
    RotateTransition, get_transition
)

# Set up logging
logging.basicConfig(level=logging.INFO,
                   format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def create_transition_video(first_video: str, second_video: str, output_path: str,
                          transition_type: str, duration: float = 2.0, **kwargs):
    """
    Create a video with a transition between two input videos.

    Args:
        first_video: Path to first video
        second_video: Path to second video
        output_path: Path to save output video
        transition_type: Type of transition ('fade', 'dissolve', 'wipe', etc.)
        duration: Transition duration in seconds
        **kwargs: Additional parameters for the transition
    """
    logger.info(f"Creating {transition_type} transition from {first_video} to {second_video}")

    # Open input videos
    cap1 = cv2.VideoCapture(first_video)
    cap2 = cv2.VideoCapture(second_video)

    if not cap1.isOpened():
        logger.error(f"Could not open first video: {first_video}")
        return

    if not cap2.isOpened():
        logger.error(f"Could not open second video: {second_video}")
        return

    # Get video properties
    width = int(cap1.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap1.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap1.get(cv2.CAP_PROP_FPS)

    # Ensure second video has same dimensions
    width2 = int(cap2.get(cv2.CAP_PROP_FRAME_WIDTH))
    height2 = int(cap2.get(cv2.CAP_PROP_FRAME_HEIGHT))

    if width != width2 or height != height2:
        logger.warning("Videos have different dimensions. Resizing second video.")

    # Calculate number of frames for each part
    # We'll use 5 seconds from each video
    video_duration = 5.0  # seconds
    video_frames = int(video_duration * fps)
    transition_frames = int(duration * fps)

    # Create output video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    # Create transition object
    transition = get_transition(transition_type, duration=duration, **kwargs)

    # Part 1: First video
    logger.info(f"Processing first video segment ({video_frames} frames)")
    for _ in range(video_frames):
        ret, frame = cap1.read()
        if not ret:
            # If we reach the end, repeat the last frame
            cap1.set(cv2.CAP_PROP_POS_FRAMES, cap1.get(cv2.CAP_PROP_POS_FRAMES) - 1)
            ret, frame = cap1.read()

        # Resize if needed
        if frame.shape[1] != width or frame.shape[0] != height:
            frame = cv2.resize(frame, (width, height))

        out.write(frame)

    # Part 2: Transition
    logger.info(f"Processing transition ({transition_frames} frames)")

    # Store the last frame of the first video
    cap1.set(cv2.CAP_PROP_POS_FRAMES, cap1.get(cv2.CAP_PROP_POS_FRAMES) - 1)
    ret1, last_frame1 = cap1.read()

    # Get the first frame of the second video
    cap2.set(cv2.CAP_PROP_POS_FRAMES, 0)
    ret2, first_frame2 = cap2.read()

    # Resize if needed
    if last_frame1.shape[1] != width or last_frame1.shape[0] != height:
        last_frame1 = cv2.resize(last_frame1, (width, height))

    if first_frame2.shape[1] != width or first_frame2.shape[0] != height:
        first_frame2 = cv2.resize(first_frame2, (width, height))

    # Apply transition
    for i in range(transition_frames):
        progress = i / (transition_frames - 1) if transition_frames > 1 else 1.0

        # Apply transition effect
        result = transition.apply(last_frame1, first_frame2, progress)

        # Write frame
        out.write(result)

    # Part 3: Second video
    logger.info(f"Processing second video segment ({video_frames} frames)")
    for _ in range(video_frames):
        ret, frame = cap2.read()
        if not ret:
            # If we reach the end, repeat the last frame
            cap2.set(cv2.CAP_PROP_POS_FRAMES, cap2.get(cv2.CAP_PROP_POS_FRAMES) - 1)
            ret, frame = cap2.read()

        # Resize if needed
        if frame.shape[1] != width or frame.shape[0] != height:
            frame = cv2.resize(frame, (width, height))

        out.write(frame)

    # Clean up
    cap1.release()
    cap2.release()
    out.release()

    logger.info(f"Transition video saved to {output_path}")

def create_transition_demo(first_video: str, second_video: str, output_dir: str):
    """
    Create a demo showcasing all available transitions.

    Args:
        first_video: Path to first video
        second_video: Path to second video
        output_dir: Directory to save output videos
    """
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)

    # Define transitions to demonstrate
    transitions = [
        ("fade", {"fade_to_black": True}),
        ("fade", {"fade_to_black": False, "duration": 1.5}),
        ("dissolve", {"mode": "simple"}),
        ("dissolve", {"mode": "zoom"}),
        ("dissolve", {"mode": "blur"}),
        ("wipe", {"direction": "left-to-right"}),
        ("wipe", {"direction": "right-to-left"}),
        ("wipe", {"direction": "top-to-bottom"}),
        ("wipe", {"direction": "bottom-to-top"}),
        ("slide", {"direction": "left-to-right"}),
        ("slide", {"direction": "right-to-left"}),
        ("pixelate", {"max_pixel_size": 30}),
        ("zoom", {"zoom_factor": 2.5}),
        ("rotate", {"direction": "horizontal"}),
        ("rotate", {"direction": "vertical"})
    ]

    logger.info(f"Creating transition demos in {output_dir}")

    # Create each transition video
    for i, (transition_type, kwargs) in enumerate(transitions):
        # Create descriptive filename
        filename = f"{i+1:02d}_{transition_type}"

        # Add parameter details to filename
        for key, value in kwargs.items():
            if key != "duration":  # Skip duration in filename
                filename += f"_{key}_{value}"

        output_path = os.path.join(output_dir, f"{filename}.mp4")

        # Create transition video
        create_transition_video(first_video, second_video, output_path,
                              transition_type, **kwargs)

def create_transition_montage(transition_dir: str, output_path: str):
    """
    Create a montage video of all transitions.

    Args:
        transition_dir: Directory containing transition videos
        output_path: Path to save montage video
    """
    # Find all transition videos
    video_files = sorted([f for f in os.listdir(transition_dir)
                        if f.endswith('.mp4')])

    if not video_files:
        logger.error(f"No transition videos found in {transition_dir}")
        return

    logger.info(f"Creating montage of {len(video_files)} transitions")

    # Get dimensions from first video
    first_video = os.path.join(transition_dir, video_files[0])
    cap = cv2.VideoCapture(first_video)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    cap.release()

    # Calculate grid dimensions
    grid_cols = min(4, len(video_files))
    grid_rows = (len(video_files) + grid_cols - 1) // grid_cols

    # Calculate montage dimensions
    montage_width = grid_cols * width
    montage_height = grid_rows * height

    # Create video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (montage_width, montage_height))

    # Open all videos
    caps = []
    for video_file in video_files:
        cap = cv2.VideoCapture(os.path.join(transition_dir, video_file))
        caps.append(cap)

    # Get video length (use shortest video)
    frame_counts = [int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) for cap in caps]
    total_frames = min(frame_counts)

    # Process frames
    for frame_idx in range(total_frames):
        # Create montage frame
        montage = np.zeros((montage_height, montage_width, 3), dtype=np.uint8)

        # Read frame from each video
        for i, cap in enumerate(caps):
            ret, frame = cap.read()
            if not ret:
                continue

            # Calculate position in grid
            row = i // grid_cols
            col = i % grid_cols

            # Copy frame to montage
            y = row * height
            x = col * width
            montage[y:y+height, x:x+width] = frame

            # Add label
            label = video_files[i].split('.')[0]
            cv2.putText(montage, label, (x + 10, y + 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

        # Write montage frame
        out.write(montage)

        # Progress update
        if frame_idx % 30 == 0:
            logger.info(f"Processing montage frame {frame_idx}/{total_frames}")

    # Clean up
    for cap in caps:
        cap.release()
    out.release()

    logger.info(f"Transition montage saved to {output_path}")

def main():
    """Main function to run the example."""
    # Hardcoded sample videos to try
    sample_video_paths = [
        os.path.join("test_media", "test_video.mp4"),
        os.path.join("test_media", "sample.mp4"),
        os.path.join("test_media", "sample2.mp4")
    ]

    # Find the first existing sample video
    first_video = None
    for path in sample_video_paths:
        if os.path.exists(path):
            first_video = path
            second_video = path  # Use the same video for both inputs
            break

    # If no sample videos found, create a better test video with colored content
    if first_video is None:
        logger.warning("No sample videos found. Creating a test video with moving content.")

        # Create output directory if it doesn't exist
        output_dir = "transition_output"
        os.makedirs(output_dir, exist_ok=True)

        # Create test_media directory if it doesn't exist
        os.makedirs("test_media", exist_ok=True)

        # Create a video with moving content instead of just black
        test_video_path = os.path.join("test_media", "generated_test.mp4")

        # Video properties
        width, height = 640, 360
        fps = 30.0
        duration = 3  # seconds

        # Create video writer
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(test_video_path, fourcc, fps, (width, height))

        # Generate frames with animated content
        for i in range(int(duration * fps)):
            # Create frame with gradient background
            frame = np.zeros((height, width, 3), dtype=np.uint8)

            # Add gradient background
            for y in range(height):
                for x in range(width):
                    b = int(255 * (0.5 + 0.5 * np.sin(0.1 * i + x * 0.01)))
                    g = int(255 * (0.5 + 0.5 * np.sin(0.1 * i + y * 0.01)))
                    r = int(255 * (0.5 + 0.5 * np.cos(0.1 * i + (x+y) * 0.01)))
                    frame[y, x] = [b, g, r]

            # Add some moving shapes
            # Moving circle
            cv2.circle(
                frame,
                center=(int(width/2 + np.sin(i*0.1)*width/3), int(height/2)),
                radius=30,
                color=(0, 0, 255),
                thickness=-1
            )

            # Moving rectangle
            cv2.rectangle(
                frame,
                pt1=(int(width/2 - 50 + np.cos(i*0.1)*width/3), int(height/2 + 70)),
                pt2=(int(width/2 + 50 + np.cos(i*0.1)*width/3), int(height/2 + 120)),
                color=(0, 255, 0),
                thickness=-1
            )

            # Write frame
            out.write(frame)

        out.release()
        logger.info(f"Created test video: {test_video_path}")
        first_video = test_video_path
        second_video = test_video_path

    # Set up parameters
    output_dir = "transition_output"
    os.makedirs(output_dir, exist_ok=True)

    # Create demo of all transitions
    try:
        create_transition_demo(first_video, second_video, output_dir)

        # Create montage
        montage_path = os.path.join(output_dir, "transitions_montage.mp4")
        create_transition_montage(output_dir, montage_path)

        logger.info(f"Transitions examples created in {output_dir}")
        logger.info(f"Montage created at {montage_path}")
    except Exception as e:
        logger.error(f"Error during transitions processing: {e}")
        return 1

    return 0

if __name__ == "__main__":
    sys.exit(main())
