import gradio as gr
import random
import time
from datetime import datetime
import tempfile
import os
from moviepy.editor import ImageClip, concatenate_videoclips
from gradio_client import Client
from PIL import Image
import edge_tts
import asyncio
import warnings
import numpy as np

warnings.filterwarnings('ignore')

# Initialize Gradio clients with public demo spaces
def initialize_clients():
    try:
        # Use a simpler public demo space
        image_client = Client("awacke1/Arxiv-Paper-Search-And-QA-RAG-Pattern")
        return image_client
    except Exception as e:
        print(f"Error initializing clients: {str(e)}")
        return None

if gr.NO_RELOAD:
    # Initialize client in NO_RELOAD block to prevent multiple initializations
    CLIENT = initialize_clients()

STORY_GENRES = [
    "Science Fiction",
    "Fantasy",
    "Mystery",
    "Romance",
    "Horror",
    "Adventure",
    "Historical Fiction",
    "Comedy"
]

STORY_STRUCTURES = {
    "Three Act": "Setup (Introduction, Inciting Incident) -> Confrontation (Rising Action, Climax) -> Resolution (Falling Action, Conclusion)",
    "Hero's Journey": "Ordinary World -> Call to Adventure -> Trials -> Transformation -> Return",
    "Five Act": "Exposition -> Rising Action -> Climax -> Falling Action -> Resolution",
    "Seven Point": "Hook -> Plot Turn 1 -> Pinch Point 1 -> Midpoint -> Pinch Point 2 -> Plot Turn 2 -> Resolution"
}

async def generate_speech(text, voice="en-US-AriaNeural"):
    """Generate speech from text using edge-tts"""
    try:
        communicate = edge_tts.Communicate(text, voice)
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp_file:
            tmp_path = tmp_file.name
            await communicate.save(tmp_path)
        return tmp_path
    except Exception as e:
        print(f"Error in text2speech: {str(e)}")
        return None

def generate_story_prompt(base_prompt, genre, structure):
    """Generate an expanded story prompt based on genre and structure"""
    prompt = f"""Create a {genre} story using this concept: '{base_prompt}'
    Follow this structure: {STORY_STRUCTURES[structure]}
    Include vivid descriptions and sensory details.
    Make it engaging and suitable for visualization.
    Keep each scene description clear and detailed enough for image generation.
    Limit the story to 5-7 key scenes.
    """
    return prompt

def generate_story(prompt, model_choice):
    """Generate story using specified model"""
    try:
        if CLIENT is None:
            return "Error: Story generation service is not available."
        
        result = CLIENT.predict(
            prompt,
            model_choice,
            True,
            api_name="/ask_llm"
        )
        return result
    except Exception as e:
        return f"Error generating story: {str(e)}"

def process_story(story_text, num_scenes=5):
    """Break story into scenes for visualization"""
    if not story_text:
        return []
        
    sentences = story_text.split('.')
    scenes = []
    scene_length = max(1, len(sentences) // num_scenes)
    
    for i in range(0, len(sentences), scene_length):
        scene = '. '.join(sentences[i:i+scene_length]).strip()
        if scene:
            scenes.append(scene)
    
    return scenes[:num_scenes]

def story_generator_interface(prompt, genre, structure, model_choice, num_scenes, words_per_scene):
    """Main story generation and multimedia creation function"""
    try:
        # Generate expanded prompt
        story_prompt = generate_story_prompt(prompt, genre, structure)
        
        # Generate story
        story = generate_story(story_prompt, model_choice)
        if story.startswith("Error"):
            return story, None, None, None
        
        # Generate speech
        audio_path = asyncio.run(generate_speech(story))
        
        return story, None, audio_path, None
        
    except Exception as e:
        error_msg = f"An error occurred: {str(e)}"
        return error_msg, None, None, None

# Create Gradio interface
with gr.Blocks(title="AI Story Generator") as demo:
    gr.Markdown("# 🎭 AI Story Generator")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Story Concept",
                placeholder="Enter your story idea...",
                lines=3
            )
            genre_input = gr.Dropdown(
                label="Genre",
                choices=STORY_GENRES,
                value="Fantasy"
            )
            structure_input = gr.Dropdown(
                label="Story Structure",
                choices=list(STORY_STRUCTURES.keys()),
                value="Three Act"
            )
            model_choice = gr.Dropdown(
                label="Model",
                choices=["mistralai/Mixtral-8x7B-Instruct-v0.1", "mistralai/Mistral-7B-Instruct-v0.2"],
                value="mistralai/Mixtral-8x7B-Instruct-v0.1"
            )
            num_scenes = gr.Slider(
                label="Number of Scenes",
                minimum=3,
                maximum=7,
                value=5,
                step=1
            )
            words_per_scene = gr.Slider(
                label="Words per Scene",
                minimum=20,
                maximum=100,
                value=50,
                step=10
            )
            generate_btn = gr.Button("Generate Story")
    
    with gr.Row():
        with gr.Column():
            story_output = gr.Textbox(
                label="Generated Story",
                lines=10,
                interactive=False  # Changed from readonly to interactive=False
            )
    
    with gr.Row():
        audio_output = gr.Audio(label="Story Narration")
    
    generate_btn.click(
        fn=story_generator_interface,
        inputs=[prompt_input, genre_input, structure_input, model_choice, num_scenes, words_per_scene],
        outputs=[story_output, None, audio_output, None]  # Set image and video outputs to None for now
    )

if __name__ == "__main__":
    demo.launch(reload=True)