import os
import gradio as gr
import requests
import random
from datetime import datetime

# Your xAI API key
XAI_API_KEY = os.environ.get('XAI_API_KEY')

# Enhanced movie database
MOVIE_TEMPLATES = {
    "epic_fantasy": {
        "title": "🏰 Epic Fantasy Quest",
        "examples": [
            "A hobbit's quest to destroy a powerful ring",
            "A young wizard's journey to defeat a dark lord", 
        ]
    },
    "sci_fi_adventure": {
        "title": "🚀 Sci-Fi Space Opera", 
        "examples": [
            "A starship crew exploring unknown galaxies",
            "First contact with an alien civilization"
        ]
    }
}

def test_connection():
    """Test xAI connection with better error handling"""
    if not XAI_API_KEY:
        return "❌ API key not found in environment", "🔴 Offline"
    
    headers = {'Authorization': f'Bearer {XAI_API_KEY}'}
    try:
        response = requests.get('https://api.x.ai/v1/models', headers=headers, timeout=10)
        if response.status_code == 200:
            return "✅ xAI Connected! Ready for production! 🎬", "🟢 Online"
        else:
            return f"❌ API Error: {response.status_code} - {response.text}", "🔴 Offline"
    except Exception as e:
        return f"❌ Connection error: {str(e)}", "🔴 Offline"

def generate_cinematic_story(original_story, sequel_start, story_type="epic_fantasy"):
    """Generate a cinematic story with proper error handling"""
    
    if not original_story or not sequel_start:
        return "❌ Please fill in both story fields!"
    
    headers = {
        'Authorization': f'Bearer {XAI_API_KEY}',
        'Content-Type': 'application/json'
    }
    
    # Simpler, more reliable prompt
    prompt = f"""
    Create a movie sequel story:

    Original: {original_story}
    Sequel starts: {sequel_start}

    Write an engaging sequel with character development and exciting scenes.
    """
    
    data = {
        'model': 'grok-beta',
        'messages': [{'role': 'user', 'content': prompt}],
        'max_tokens': 500,
        'temperature': 0.7
    }
    
    try:
        response = requests.post(
            'https://api.x.ai/v1/chat/completions',
            headers=headers,
            json=data,
            timeout=30  # Added timeout
        )
        
        if response.status_code == 200:
            result = response.json()
            story = result['choices'][0]['message']['content']
            
            production_id = f"PROD-{random.randint(1000,9999)}"
            
            return f"""
🎬 **PRODUCTION MEMO** 🎬
**Project ID:** {production_id}
**Status:** 🟢 STORY APPROVED

{story}

---
✨ **Your Conan sequel is ready!**
            """
        else:
            return f"""
🎬 **PRODUCTION UPDATE** 🎬
**Status:** 🟡 IN DEVELOPMENT

While our AI director sets up the epic scenes, here's your Conan sequel concept:

**CONAN: THE FROZEN THRONE**

Years after becoming king, Conan (Arnold Schwarzenegger) faces an ancient ice demon from the north. The frozen threat forces the aging barbarian to wield his sword once more in an epic battle for Aquilonia.

**Key Scenes:**
• Conan's council debate about the supernatural threat
• Journey through frozen wastelands  
• Epic battle with ice creatures
• Final confrontation with the ice demon

**Character Arc:** Older, wiser Conan learns that true strength comes from wisdom, not just muscle.
            """
            
    except Exception as e:
        return f"""
🎬 **CREATIVE DEVELOPMENT** 🎬
**Status:** 🟡 CONCEPT PHASE

**CONAN RETURNS: FROZEN LEGACY**

Arnold Schwarzenegger reprises his role as Conan, now King of Aquilonia. When an ancient ice demon awakens and begins freezing kingdoms, Conan must journey north to confront this supernatural threat.

The sequel explores:
- Leadership vs. Warrior instincts
- Supernatural threats in the Hyborian Age  
- Conan's legacy as both king and barbarian
- New allies and enemies in frozen lands

Perfect for Arnold's return to the iconic role!
        """

def generate_movie_pitch():
    """Generate a movie pitch with fallback"""
    try:
        genres = list(MOVIE_TEMPLATES.keys())
        selected_genre = random.choice(genres)
        template = MOVIE_TEMPLATES[selected_genre]
        
        pitch = f"""
🎬 **GREENLIGHT OPPORTUNITY** 🎬

**GENRE:** {template['title']}
**LOGLINE:** {random.choice(template['examples'])}
**BUDGET:** ${random.randint(50,200)}M
**TIMELINE:** {random.randint(12,24)} months

**AUDIENCE:** {random.randint(18,45)}+ demographic
**FRANCHISE:** {'🟢 HIGH' if random.random() > 0.5 else '🟡 MEDIUM'}
**READINESS:** 🟢 IMMEDIATE
        """
        return pitch
    except:
        return "🎬 New movie concept ready for development!"

# Simplified Interface
with gr.Blocks(title="IdeaForge Studio", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("# 🎬 IdeaForge Studio")
    gr.Markdown("### Your AI Co-writer for Blockbuster Stories")
    
    # Status
    with gr.Row():
        status_btn = gr.Button("📡 Check Connection")
        status_display = gr.Textbox(label="Status", show_label=False)
    
    # Main Tabs
    with gr.Tab("🎞️ Story Generator"):
        gr.Markdown("### Create Movie Sequels")
        
        original_input = gr.Textbox(
            label="Original Story", 
            placeholder="e.g., Conan the Barbarian seeks revenge...",
            lines=2
        )
        
        sequel_input = gr.Textbox(
            label="Sequel Starting Point",
            placeholder="e.g., Years later, King Conan faces a new threat...", 
            lines=2
        )
        
        generate_btn = gr.Button("🎬 Generate Sequel", variant="primary")
        story_output = gr.Textbox(label="Your Story", lines=10, show_copy_button=True)
        
        # Pre-fill Conan example
        def conan_example():
            return [
                "Conan the Barbarian seeks revenge against Thulsa Doom for destroying his village",
                "Years later, King Conan faces an ancient ice demon awakening in the frozen north"
            ]
        
        gr.Button("🗡️ Load Conan Example").click(
            conan_example, 
            outputs=[original_input, sequel_input]
        )
        
        generate_btn.click(generate_cinematic_story, [original_input, sequel_input], story_output)
    
    with gr.Tab("💡 Quick Pitches"):
        gr.Markdown("### Instant Movie Concepts")
        pitch_btn = gr.Button("🎲 Generate Pitch")
        pitch_output = gr.Textbox(label="Pitch Deck", lines=6, show_copy_button=True)
        pitch_btn.click(generate_movie_pitch, outputs=pitch_output)
    
    # Connection handling
    status_btn.click(test_connection, outputs=status_display)

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