import gradio as gr
import random
from datetime import datetime
import tempfile
import os
import edge_tts
import asyncio
import warnings
from gradio_client import Client
import pytz
import re
import json

warnings.filterwarnings('ignore')

# Initialize client outside of interface definition
arxiv_client = None

def init_client():
    global arxiv_client
    if arxiv_client is None:
        arxiv_client = Client("awacke1/Arxiv-Paper-Search-And-QA-RAG-Pattern")
    return arxiv_client

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

async def generate_speech(text, voice="en-US-AriaNeural"):
    """Generate speech from text"""
    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 process_story_and_audio(prompt, model_choice):
    """Process story and generate audio"""
    try:
        # Generate story
        story = generate_story(prompt, model_choice)
        if isinstance(story, str) and story.startswith("Error"):
            return story, None

        # Generate audio
        audio_path = asyncio.run(generate_speech(story))
        
        return story, audio_path
    except Exception as e:
        return f"Error: {str(e)}", None

# Create the Gradio interface
with gr.Blocks(title="AI Story Generator") as demo:
    gr.Markdown("""
    # 🎭 AI Story Generator & Narrator
    Generate creative stories and listen to them!
    """)
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Story Concept",
                placeholder="Enter your story idea...",
                lines=3
            )
            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"
            )
            generate_btn = gr.Button("Generate Story")
    
    with gr.Row():
        story_output = gr.Textbox(
            label="Generated Story",
            lines=10,
            interactive=False
        )
    
    with gr.Row():
        audio_output = gr.Audio(
            label="Story Narration",
            type="filepath"
        )
    
    generate_btn.click(
        fn=process_story_and_audio,
        inputs=[prompt_input, model_choice],
        outputs=[story_output, audio_output]
    )

# Launch the app using the current pattern
if __name__ == "__main__":
    demo.launch()