import streamlit as st
import torch
from transformers import pipeline, AutoProcessor, AutoModel
from PIL import Image
import numpy as np
import cv2
import tempfile
import os

def save_video_frames(frames, output_path, fps=30):
    """
    Save video frames using OpenCV instead of moviepy
    """
    # Get frame dimensions
    height, width = frames[0].shape[:2]
    
    # Initialize video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    # Write frames
    for frame in frames:
        # Convert from RGB to BGR (OpenCV uses BGR)
        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        out.write(frame_bgr)
    
    # Release video writer
    out.release()

def generate_video_from_image(image, duration_seconds=10, progress_bar=None):
    """
    Generate a video from an image using LTX-Video and image captioning.
    """
    try:
        if progress_bar:
            progress_bar.progress(0.1, "Generating image caption...")
            
        # Setup image captioning pipeline
        captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
        
        # Generate caption
        caption = captioner(image)[0]['generated_text']
        st.write(f"Generated caption: *{caption}*")
        
        if progress_bar:
            progress_bar.progress(0.3, "Loading LTX-Video model...")
            
        # Initialize LTX-Video components
        processor = AutoProcessor.from_pretrained("Lightricks/ltx-video")
        model = AutoModel.from_pretrained("Lightricks/ltx-video")
        
        if progress_bar:
            progress_bar.progress(0.4, "Processing image...")
            
        # Process image
        inputs = processor(images=image, return_tensors="pt")
        
        if progress_bar:
            progress_bar.progress(0.5, "Generating video frames...")
            
        # Generate video frames
        num_frames = duration_seconds * 30  # 30 FPS
        with torch.no_grad():
            video_frames = model.generate(
                **inputs,
                num_frames=num_frames,
                num_inference_steps=50,
                guidance_scale=7.5,
                prompt=caption,
            ).videos[0]
        
        if progress_bar:
            progress_bar.progress(0.8, "Creating final video...")
            
        # Convert frames to numpy arrays
        frames = [np.array(frame) for frame in video_frames]
        
        # Create temporary file for video
        with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_file:
            output_path = tmp_file.name
            
        # Save video using OpenCV
        save_video_frames(frames, output_path)
        
        if progress_bar:
            progress_bar.progress(1.0, "Video generation complete!")
            
        return output_path, caption
        
    except Exception as e:
        st.error(f"Error generating video: {str(e)}")
        raise  # Re-raise the exception to see the full error message
        return None, None

def main():
    st.set_page_config(page_title="Video Generator", page_icon="🎥")
    
    st.title("🎥 AI Video Generator")
    st.write("""
    Upload an image to generate a video with AI-powered motion and transitions.
    The app will automatically generate a caption for your image and use it as inspiration for the video.
    """)
    
    # File uploader
    uploaded_file = st.file_uploader("Choose an image", type=['png', 'jpg', 'jpeg'])
    
    # Duration selector
    duration = st.slider("Video duration (seconds)", min_value=1, max_value=30, value=10)
    
    if uploaded_file is not None:
        # Display uploaded image
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image", use_column_width=True)
        
        # Generate button
        if st.button("Generate Video"):
            try:
                # Create a progress bar
                progress_text = "Operation in progress. Please wait..."
                my_bar = st.progress(0, text=progress_text)
                
                # Generate video
                video_path, caption = generate_video_from_image(image, duration, my_bar)
                
                if video_path and os.path.exists(video_path):
                    # Read the video file
                    with open(video_path, 'rb') as video_file:
                        video_bytes = video_file.read()
                    
                    # Create download button
                    st.download_button(
                        label="Download Video",
                        data=video_bytes,
                        file_name="generated_video.mp4",
                        mime="video/mp4"
                    )
                    
                    # Display video
                    st.video(video_bytes)
                    
                    # Clean up temporary file
                    os.unlink(video_path)
                else:
                    st.error("Failed to generate video. Please try again.")
            
            except Exception as e:
                st.error(f"An error occurred: {str(e)}")

if __name__ == "__main__":
    main()