import os
import gradio as gr
from groq import Groq

# Initialize Groq client with API key from environment
GROQ_API_KEY = os.getenv("GROQ_API_KEY", "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")
client = Groq(api_key=GROQ_API_KEY)

# Store conversation history
conversation_history = []
max_history_length = 10

def stream_response(messages, temperature=1):
    """Stream response from Groq API."""
    completion = client.chat.completions.create(
        model="openai/gpt-oss-120b",
        messages=messages,
        temperature=temperature,
        max_completion_tokens=8192,
        top_p=1,
        stream=True,
        stop=None
    )
    
    response_text = ""
    for chunk in completion:
        if chunk.choices[0].delta.content:
            response_text += chunk.choices[0].delta.content
            yield response_text

def generate_story(story_prompt, story_type):
    """Generate a story based on prompt and type."""
    global conversation_history
    
    if not story_prompt.strip():
        return "Please enter a story prompt."
    
    # Create system message based on story type
    system_messages = {
        "Fantasy": "You are a creative fantasy storyteller. Generate an engaging, immersive fantasy story with vivid descriptions.",
        "Science Fiction": "You are a sci-fi storyteller. Generate an engaging science fiction story with futuristic elements and world-building.",
        "Mystery": "You are a mystery writer. Generate a compelling mystery story with plot twists and clues.",
        "Romance": "You are a romance novelist. Generate a heartwarming, emotional romance story.",
        "Horror": "You are a horror writer. Generate a suspenseful, atmospheric horror story.",
        "Adventure": "You are an adventure writer. Generate an exciting adventure story with action and discovery.",
        "General": "You are a creative writer. Generate an engaging and well-structured story."
    }
    
    system_prompt = system_messages.get(story_type, system_messages["General"])
    
    # Build messages for API
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"Write a story based on this prompt: {story_prompt}"}
    ]
    
    # Add conversation history for context
    if conversation_history:
        messages = conversation_history[-max_history_length:] + messages
    
    # Stream and collect response
    full_response = ""
    for partial_response in stream_response(messages):
        full_response = partial_response
        yield full_response
    
    # Store in conversation history
    conversation_history.append({"role": "user", "content": f"Story ({story_type}): {story_prompt}"})
    conversation_history.append({"role": "assistant", "content": full_response})

def generate_storyboard(scenario):
    """Generate a detailed storyboard for a given scenario."""
    global conversation_history
    
    if not scenario.strip():
        return "Please enter a scenario for the storyboard."
    
    storyboard_prompt = f"""Generate a detailed storyboard for the following scenario:

Scenario: {scenario}

Please create a storyboard with:
1. Scene-by-scene breakdown
2. Key visual elements for each scene
3. Dialogue or narration
4. Transitions between scenes
5. Camera angles and movements
6. Lighting and mood notes
7. Props and set design

Format the output in a clear, structured manner with each scene clearly numbered and separated."""
    
    messages = [
        {"role": "system", "content": "You are a professional storyboard artist and screenwriter. Create detailed, visual storyboards with cinematic descriptions."},
        {"role": "user", "content": storyboard_prompt}
    ]
    
    full_response = ""
    for partial_response in stream_response(messages):
        full_response = partial_response
        yield full_response
    
    # Store in history
    conversation_history.append({"role": "user", "content": f"Storyboard: {scenario}"})
    conversation_history.append({"role": "assistant", "content": full_response})

def chat(user_message, chat_history):
    """Handle general chat conversations."""
    global conversation_history
    
    if not user_message.strip():
        return chat_history
    
    # Add user message to chat history
    chat_history.append((user_message, ""))
    
    # Build messages for API
    messages = [
        {"role": "system", "content": "You are a helpful AI assistant specializing in creative writing, storytelling, and narrative development."}
    ]
    
    # Add chat history to context (limit to prevent token overflow)
    for user_msg, assistant_msg in chat_history[-max_history_length:]:
        messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    
    # Stream response
    full_response = ""
    for partial_response in stream_response(messages):
        full_response = partial_response
        chat_history[-1] = (user_message, full_response)
        yield chat_history
    
    # Store in conversation history
    conversation_history.append({"role": "user", "content": user_message})
    conversation_history.append({"role": "assistant", "content": full_response})

def reset_chat_history():
    """Reset chat history."""
    return []

def reset_all_history():
    """Reset all conversation history."""
    global conversation_history
    conversation_history = []
    return []

# Create Gradio interface
def create_interface():
    with gr.Blocks(
        title="Story Generator & Chatbot",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1200px;
            margin: auto;
        }
        """
    ) as demo:
        gr.Markdown(
            """
            # 📖 AI Story Generator & Chatbot
            
            **Powered by Groq API** - Generate creative stories, create storyboards, and chat about narrative writing!
            """
        )
        
        with gr.Tabs():
            # Story Generation Tab
            with gr.Tab("🎭 Story Generator"):
                gr.Markdown("### Generate creative stories")
                gr.Markdown("Provide a story prompt and select a genre to generate a unique story.")
                
                with gr.Row():
                    with gr.Column(scale=2):
                        story_prompt = gr.Textbox(
                            label="Story Prompt",
                            placeholder="e.g., A mysterious letter arrives at an old mansion on a stormy night...",
                            lines=4,
                            info="Be descriptive for better results"
                        )
                        
                        with gr.Row():
                            story_type = gr.Dropdown(
                                choices=["Fantasy", "Science Fiction", "Mystery", "Romance", "Horror", "Adventure", "General"],
                                value="General",
                                label="Story Type",
                                info="Choose the genre for your story"
                            )
                            generate_btn = gr.Button("✨ Generate Story", variant="primary", scale=1)
                    
                    with gr.Column(scale=2):
                        story_output = gr.Textbox(
                            label="Generated Story",
                            lines=20,
                            interactive=False,
                            show_copy_button=True
                        )
                
                generate_btn.click(
                    generate_story,
                    inputs=[story_prompt, story_type],
                    outputs=story_output
                )
            
            # Storyboard Tab
            with gr.Tab("🎬 Storyboard Creator"):
                gr.Markdown("### Create detailed storyboards")
                gr.Markdown("Describe your scenario in detail to get a comprehensive visual storyboard.")
                
                with gr.Row():
                    with gr.Column(scale=1):
                        scenario = gr.Textbox(
                            label="Scenario Description",
                            placeholder="e.g., A superhero saves the city from an alien invasion. Include setting, characters, and key events...",
                            lines=5,
                            info="More details = better storyboard"
                        )
                        storyboard_btn = gr.Button("🎥 Generate Storyboard", variant="primary")
                    
                    with gr.Column(scale=1):
                        storyboard_output = gr.Textbox(
                            label="Generated Storyboard",
                            lines=20,
                            interactive=False,
                            show_copy_button=True
                        )
                
                storyboard_btn.click(
                    generate_storyboard,
                    inputs=scenario,
                    outputs=storyboard_output
                )
            
            # Chat Tab
            with gr.Tab("💬 Storytelling Chat"):
                gr.Markdown("### Chat about creative writing")
                gr.Markdown("Discuss storytelling techniques, get writing advice, and brainstorm ideas.")
                
                with gr.Row():
                    with gr.Column():
                        chatbot = gr.Chatbot(
                            label="Conversation",
                            height=500,
                            show_label=True,
                            scale=1
                        )
                
                with gr.Row():
                    user_input = gr.Textbox(
                        label="Your message",
                        placeholder="Ask about storytelling, writing tips, or brainstorm ideas...",
                        scale=5,
                        show_label=False
                    )
                    send_btn = gr.Button("Send", variant="primary", scale=1)
                
                with gr.Row():
                    clear_chat_btn = gr.Button("Clear Chat", scale=1)
                
                # Handle sending messages
                def send_message(msg, history):
                    return chat(msg, history)
                
                send_btn.click(
                    send_message,
                    inputs=[user_input, chatbot],
                    outputs=chatbot
                ).then(
                    lambda: "",
                    outputs=user_input
                )
                
                user_input.submit(
                    send_message,
                    inputs=[user_input, chatbot],
                    outputs=chatbot
                ).then(
                    lambda: "",
                    outputs=user_input
                )
                
                clear_chat_btn.click(
                    reset_chat_history,
                    outputs=chatbot
                )
        
        gr.Markdown(
            """
            ---
            **Tips for best results:**
            - 📝 Be detailed and descriptive in your prompts
            - 🎯 The more context you provide, the better the output
            - 💡 Use story types to guide the tone and genre
            - 🔄 Refine your ideas through conversation
            
            **Powered by:** [Groq API](https://groq.com) | **Built with:** [Gradio](https://gradio.app)
            """
        )
    
    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        max_threads=10
    )