import os
os.environ['MPLCONFIGDIR'] = '/tmp/matplotlib'
os.environ['FONTCONFIG_PATH'] = '/tmp/fontconfig'

import torch
import gradio as gr
import tempfile
import logging
from models import load_models
from video_generator import generate_video_pipeline

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

# Hardware configuration
device = "cuda" if torch.cuda.is_available() else "cpu"
logger.info(f"Running on {device.upper()}")

# Load models
try:
    text_to_image, image_to_video, tts_model = load_models()
except Exception as e:
    logger.error(f"Initialization failed: {str(e)}")
    raise

def generate_video(prompt, duration=5, fps=24):
    with tempfile.TemporaryDirectory() as tmpdir:
        try:
            return generate_video_pipeline(
                prompt=prompt,
                text_to_image_model=text_to_image,
                image_to_video_model=image_to_video,
                tts_model=tts_model,
                output_dir=tmpdir,
                duration=duration,
                fps=fps
            )
        except Exception as e:
            logger.error(f"Generation failed: {str(e)}")
            raise gr.Error(f"Error: {str(e)}")

with gr.Blocks(title="AI Video Generator") as app:
    gr.Markdown("# 🎥 AI Video Generator")
    
    with gr.Row():
        prompt_input = gr.Textbox(label="Prompt", placeholder="A cat in space...")
    
    with gr.Row():
        duration = gr.Slider(2, 30, 5, label="Duration (s)")
        fps = gr.Slider(12, 60, 24, label="FPS")
    
    generate_btn = gr.Button("Generate", variant="primary")
    
    with gr.Row():
        video_output = gr.Video(label="Result", format="mp4")
        download_btn = gr.File(label="Download", type="file")
    
    generate_btn.click(
        generate_video,
        inputs=[prompt_input, duration, fps],
        outputs=[video_output, download_btn]
    )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0")