"""
PROFESSIONAL VIDEO GENERATOR (LOCAL & OPEN SOURCE)
====================================================
Features:
✅ 100% FREE & Open Source (runs on your hardware)
✅ State-of-the-art Stable Video Diffusion model
✅ High-quality, coherent video generation
✅ Full privacy and control, no external APIs
✅ Optimized for consumer-grade GPUs (>= 12GB VRAM recommended)
✅ Production-ready local inference engine

Author: Professional AI Solutions Architect
Date: 2024
"""

import torch
from diffusers import StableVideoDiffusionPipeline, StableDiffusionXLPipeline
from diffusers.utils import load_image, export_to_video
from PIL import Image
import time
from pathlib import Path
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ProfessionalVideoGenerator:
    """
    Video generation using a local, self-hosted Stable Video Diffusion model.
    This approach provides maximum control, privacy, and eliminates API dependencies.
    """

    def __init__(self):
        """
        Initializes the video generation pipelines.
        
        This is a heavy operation and should only be done once at application startup.
        It will download models from Hugging Face on the first run.
        """
        if not torch.cuda.is_available():
            raise RuntimeError("❌ CUDA is not available. A GPU is required for this video generator.")

        self.device = "cuda"
        self.output_dir = Path("generated_videos")
        self.output_dir.mkdir(exist_ok=True)
        self.generation_count = 0

        logger.info(" M odel Loading: This may take a few minutes on the first run...")

        try:
            # --- Load Image-to-Video Pipeline (Stable Video Diffusion) ---
            # Using fp16 for reduced VRAM usage and faster inference
            self.svd_pipe = StableVideoDiffusionPipeline.from_pretrained(
                "stabilityai/stable-video-diffusion-img2vid-xt",
                torch_dtype=torch.float16,
                variant="fp16"
            )
            # Optimization: Enable model CPU offloading to run on lower VRAM GPUs
            self.svd_pipe.enable_model_cpu_offload()

            # --- Load Text-to-Image Pipeline (SDXL) for generating the initial frame ---
            self.txt2img_pipe = StableDiffusionXLPipeline.from_pretrained(
                "stabilityai/sdxl-turbo", 
                torch_dtype=torch.float16, 
                variant="fp16"
            )
            self.txt2img_pipe = self.txt2img_pipe.to(self.device)


            logger.info("✅ ProfessionalVideoGenerator initialized (Local SVD)")
            logger.info(f"   GPU: {torch.cuda.get_device_name(0)}")
            logger.info("   Models: SDXL-Turbo (for initial image) + Stable Video Diffusion (for animation)")
            logger.info(f"   Output directory: {self.output_dir}/")

        except Exception as e:
            logger.error(f"❌ Failed to load models. Ensure you have a stable internet connection for the first download and sufficient VRAM. Error: {e}")
            raise

    def generate_video(self, prompt: str, style: str, aspect_ratio: str = "16:9") -> dict:
        """
        Generates a video from a text prompt.
        
        Workflow:
        1. Generate a high-quality initial image from the prompt using SDXL Turbo.
        2. Animate the image into a video using Stable Video Diffusion.
        
        Args:
            prompt (str): Text description for the video.
            style (str): Visual style (e.g., 'realistic', 'cinematic', 'animated').
            aspect_ratio (str): The desired aspect ratio. Note: SVD is optimized for 16:9.

        Returns:
            dict: Result with video path and metadata, or an error message.
        """
        self.generation_count += 1
        start_time = time.time()
        
        logger.info(f"\n🎬 Generating Video #{self.generation_count}")
        logger.info(f"   Prompt: {prompt[:80]}...")
        logger.info(f"   Style: {style}")

        try:
            # --- Step 1: Generate Initial Image using SDXL Turbo ---
            logger.info("   ⏳ 1/3: Generating initial image...")
            # Append style to prompt for better stylistic control
            full_prompt = f"{prompt}, {style}"
            
            # SVD works best with a specific resolution, we generate close to it
            # Native resolution for SVD is 1024x576 (16:9)
            image_width, image_height = 1024, 576
            
            initial_image = self.txt2img_pipe(
                prompt=full_prompt, 
                num_inference_steps=2, # SDXL-Turbo is fast
                guidance_scale=0.0,
                width=image_width,
                height=image_height
            ).images[0]
            
            # --- Step 2: Animate Image with Stable Video Diffusion ---
            logger.info("   ⏳ 2/3: Animating frames with SVD...")
            # The SVD model expects a specific resolution
            image = initial_image.resize((1024, 576))

            frames = self.svd_pipe(
                image, 
                decode_chunk_size=8, # Lower for less VRAM, higher for faster decoding
                motion_bucket_id=127,
                noise_aug_strength=0.02,
                num_frames=25 # Number of frames in the output video
            ).frames[0]
            
            # --- Step 3: Save Video ---
            logger.info("   ⏳ 3/3: Exporting video file...")
            timestamp = int(time.time() * 1000)
            video_filename = f"video_{timestamp}.mp4"
            video_path = self.output_dir / video_filename

            export_to_video(frames, str(video_path), fps=7)
            
            generation_time = time.time() - start_time
            video_size_mb = video_path.stat().st_size / (1024 * 1024)

            logger.info("   ✅ SUCCESS")
            logger.info(f"   ⏱️  Time: {generation_time:.1f}s")
            logger.info(f"   📊 Size: {video_size_mb:.2f} MB")
            logger.info(f"   💾 Saved: {video_path}")

            return {
                "success": True,
                "video_path": str(video_path),
                "video_filename": video_filename,
                "prompt": prompt,
                "style": style,
                "generation_time": generation_time
            }

        except Exception as e:
            import traceback
            error_msg = f"Video generation failed: {str(e)}"
            logger.error(f"   ❌ {error_msg}")
            logger.error(traceback.format_exc())
            return {"success": False, "error": error_msg}