import gradio as gr
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
import json
import requests
from typing import Dict, List, Optional
import base64
from io import BytesIO
from PIL import Image
import os

class VisualRoleplayAgent:
    def __init__(self, groq_api_key: str, stability_api_key: str = None, model_name: str = "meta-llama/llama-guard-4-12b"):
        """Initialize the visual roleplay agent"""
        self.llm = ChatGroq(
            groq_api_key=groq_api_key,
            model_name=model_name,
            temperature=0.7
        )
        self.stability_api_key = stability_api_key
        self.chat_history: List[Dict] = []
        self.current_roleplay = None
        self.character_info = {}
        self.current_scene_image = None
        
    def setup_roleplay(self, roleplay_request: str) -> Dict:
        """Parse user's roleplay request and set up the scenario"""
        
        # Create character setup
        setup_prompt = ChatPromptTemplate.from_template(
            """Based on this roleplay request: "{roleplay_request}"

Create a detailed character setup. Respond with ONLY valid JSON in this exact format:

{{
    "character_name": "Name of character I should play",
    "character_description": "Detailed description of the character",
    "setting": "Where/when the roleplay takes place", 
    "scenario": "Initial scenario/situation",
    "personality_traits": ["trait1", "trait2", "trait3"],
    "speaking_style": "How the character speaks",
    "background": "Character's background/history",
    "goals": "What the character wants to achieve"
}}

Make sure it's valid JSON with no extra text."""
        )
        
        try:
            response = self.llm.invoke(setup_prompt.format(roleplay_request=roleplay_request))
            content = response.content.strip()
            
            # Try to extract JSON
            json_start = content.find('{')
            json_end = content.rfind('}') + 1
            
            if json_start == -1 or json_end <= json_start:
                character_data = self._create_fallback_character(roleplay_request)
            else:
                json_content = content[json_start:json_end]
                try:
                    character_data = json.loads(json_content)
                except json.JSONDecodeError:
                    character_data = self._create_fallback_character(roleplay_request)
            
            self.character_info = character_data
            self.current_roleplay = roleplay_request
            self.chat_history = []
            
            return {
                "status": "success",
                "message": f"Roleplay setup complete! I'm now {self.character_info['character_name']}",
                "character_info": self.character_info
            }
            
        except Exception as e:
            character_data = self._create_fallback_character(roleplay_request)
            self.character_info = character_data
            self.current_roleplay = roleplay_request
            self.chat_history = []
            
            return {
                "status": "success",
                "message": f"Roleplay setup complete! I'm now {self.character_info['character_name']}",
                "character_info": self.character_info
            }
    
    def _create_fallback_character(self, roleplay_request: str) -> Dict:
        """Create a fallback character when JSON parsing fails"""
        request_lower = roleplay_request.lower()
        
        if 'coffee' in request_lower or 'café' in request_lower or 'barista' in request_lower:
            return {
                "character_name": "Luna 'Lulu' Thompson",
                "character_description": "Energetic barista with bright eyes and infectious enthusiasm",
                "setting": "Modern coffee shop with warm lighting and cozy atmosphere",
                "scenario": "Busy morning at the coffee shop, serving specialty drinks",
                "personality_traits": ["energetic", "friendly", "passionate", "bubbly"],
                "speaking_style": "Excitable and warm, uses lots of expressions",
                "background": "Coffee enthusiast who loves making people's day brighter",
                "goals": "Create the perfect coffee experience for every customer"
            }
        elif 'tavern' in request_lower or 'inn' in request_lower:
            return {
                "character_name": "Gareth the Innkeeper",
                "character_description": "A stout, friendly man with a graying beard and warm smile",
                "setting": "A cozy medieval tavern with flickering candles and wooden tables",
                "scenario": "Evening at the tavern, serving drinks and meals to travelers",
                "personality_traits": ["friendly", "wise", "helpful", "discreet"],
                "speaking_style": "Warm and welcoming, with occasional hints of deeper knowledge",
                "background": "Has run this tavern for decades, knows everyone in town",
                "goals": "To provide good hospitality and gather interesting news from travelers"
            }
        elif 'detective' in request_lower or 'mystery' in request_lower:
            return {
                "character_name": "Inspector Blake",
                "character_description": "A sharp-eyed detective with a keen mind and methodical approach",
                "setting": "A foggy Victorian-era city with gas lamps and cobblestone streets",
                "scenario": "Investigating a puzzling case that requires careful attention to detail",
                "personality_traits": ["observant", "logical", "persistent", "intuitive"],
                "speaking_style": "Precise and thoughtful, asks probing questions",
                "background": "Years of experience solving complex criminal cases",
                "goals": "To uncover the truth and bring justice to those who deserve it"
            }
        else:
            return {
                "character_name": "Alex the Guide",
                "character_description": "A knowledgeable and adaptable person ready to help",
                "setting": "A versatile location that can adapt to the situation",
                "scenario": "Ready to assist and engage in whatever adventure awaits",
                "personality_traits": ["helpful", "adaptable", "curious", "engaging"],
                "speaking_style": "Clear and engaging, adjusts to the situation",
                "background": "Experienced in many different situations and scenarios",
                "goals": "To provide an engaging and helpful roleplay experience"
            }
    
    def generate_roleplay_response(self, user_message: str) -> str:
        """Generate a response in character for the current roleplay"""
        
        if not self.current_roleplay:
            return "Please set up a roleplay first!"
        
        # Build conversation history
        history_str = ""
        for msg in self.chat_history[-3:]:  # Keep last 3 exchanges
            history_str += f"Human: {msg['human']}\n{self.character_info['character_name']}: {msg['ai']}\n\n"
        
        # Create roleplay prompt
        roleplay_prompt = ChatPromptTemplate.from_template(
            """You are {character_name}. Here are your details:

Character: {character_name}
Description: {character_description}
Setting: {setting}
Current Situation: {scenario}
Personality: {personality_traits}
How You Speak: {speaking_style}
Your Background: {background}
Your Goals: {goals}

Recent conversation:
{chat_history}

The human just said: "{user_message}"

Respond as {character_name}. Stay in character completely. Be natural and engaging. Don't mention being an AI."""
        )
        
        try:
            response = self.llm.invoke(roleplay_prompt.format(
                character_name=self.character_info['character_name'],
                character_description=self.character_info['character_description'],
                setting=self.character_info['setting'],
                scenario=self.character_info['scenario'],
                personality_traits=", ".join(self.character_info['personality_traits']),
                speaking_style=self.character_info['speaking_style'],
                background=self.character_info['background'],
                goals=self.character_info['goals'],
                chat_history=history_str,
                user_message=user_message
            ))
            
            # Save to history
            self.chat_history.append({
                "human": user_message,
                "ai": response.content
            })
            
            return response.content
            
        except Exception as e:
            return f"*{self.character_info['character_name']} seems momentarily confused* Sorry, could you repeat that?"
    
    def end_roleplay(self) -> str:
        """End the current roleplay session"""
        if self.current_roleplay:
            character_name = self.character_info.get('character_name', 'Character')
            self.current_roleplay = None
            self.character_info = {}
            self.chat_history = []
            return f"Roleplay session ended. {character_name} waves goodbye!"
        else:
            return "No active roleplay session to end."
        
    def generate_image_prompt(self, scene_description: str, character_name: str) -> str:
        """Generate an image prompt based on the scene"""
        prompt_template = ChatPromptTemplate.from_template(
            """Create a detailed image prompt for: {scene_description}

Make it descriptive and visual, under 100 words.
Focus on the setting, atmosphere, and mood.

Example: "Cozy medieval tavern, warm candlelight, wooden tables, stone walls, fantasy art style"

Your prompt:"""
        )
        
        try:
            response = self.llm.invoke(prompt_template.format(
                scene_description=scene_description
            ))
            return response.content.strip()
        except:
            return f"Beautiful {scene_description}, detailed, high quality"
    
    def get_placeholder_image(self, scene_type: str) -> str:
        """Get placeholder image based on scene type"""
        placeholders = {
            "coffee": "https://images.unsplash.com/photo-1501339847302-ac426a4a7cbb?w=512&h=512&fit=crop&auto=format",
            "tavern": "https://images.unsplash.com/photo-1567538096630-e0c55bd6374c?w=512&h=512&fit=crop&auto=format",
            "detective": "https://images.unsplash.com/photo-1516962126636-27ad087061cc?w=512&h=512&fit=crop&auto=format",
            "fantasy": "https://images.unsplash.com/photo-1578662996442-48f60103fc96?w=512&h=512&fit=crop&auto=format",
            "space": "https://images.unsplash.com/photo-1446776653964-20c1d3a81b06?w=512&h=512&fit=crop&auto=format",
            "default": "https://images.unsplash.com/photo-1518709268805-4e9042af2176?w=512&h=512&fit=crop&auto=format"
        }
        
        for key, url in placeholders.items():
            if key in scene_type.lower():
                return url
        return placeholders["default"]

class VisualRoleplayApp:
    def __init__(self, groq_api_key: str, stability_api_key: str = None):
        self.agent = VisualRoleplayAgent(groq_api_key, stability_api_key)
        self.session_active = False
        
    def setup_roleplay(self, roleplay_request: str):
        """Setup roleplay and generate initial scene"""
        result = self.agent.setup_roleplay(roleplay_request)
        
        if result['status'] == 'success':
            self.session_active = True
            
            # Get placeholder image based on setting
            scene_image = self.agent.get_placeholder_image(result['character_info']['setting'])
            
            return {
                "status": "success",
                "character_info": result['character_info'],
                "scene_image": scene_image,
                "opening_message": f"Welcome to {result['character_info']['setting']}!"
            }
        
        return result
    
    def create_gradio_interface(self):
        """Create Gradio interface for visual roleplay"""
        with gr.Blocks(title="Visual Roleplay System", theme=gr.themes.Soft()) as interface:
            gr.Markdown("# 🎭 Visual AI Roleplay System")
            gr.Markdown("Experience immersive roleplay with AI characters and beautiful visuals!")
            
            with gr.Row():
                with gr.Column(scale=1):
                    # Setup Panel
                    gr.Markdown("## Setup Your Roleplay")
                    roleplay_input = gr.Textbox(
                        label="Describe your roleplay scenario",
                        placeholder="e.g., I want to roleplay as a customer in a magical coffee shop...",
                        lines=3
                    )
                    setup_btn = gr.Button("🎬 Start Roleplay", variant="primary")
                    
                    # Character Info
                    with gr.Accordion("Character Information", open=False):
                        character_info = gr.JSON(label="Character Details")
                    
                with gr.Column(scale=1):
                    # Scene Visual
                    scene_image = gr.Image(
                        label="Current Scene", 
                        height=400,
                        show_label=True,
                        interactive=False
                    )
            
            with gr.Row():
                with gr.Column():
                    # Chat Interface
                    chatbot = gr.Chatbot(
                        label="Roleplay Conversation",
                        height=400,
                        avatar_images=("🧑‍💼", "🎭"),
                        show_copy_button=True
                    )
                    
                    with gr.Row():
                        user_input = gr.Textbox(
                            label="Your message",
                            placeholder="Type your message here...",
                            scale=4,
                            lines=2
                        )
                        with gr.Column(scale=1):
                            send_btn = gr.Button("📤 Send", variant="primary")
                            end_btn = gr.Button("🚪 End Roleplay", variant="secondary")
            
            # Examples
            gr.Examples(
                examples=[
                    ["I want to roleplay as a customer in a cozy coffee shop, and you are an enthusiastic barista named Luna."],
                    ["I want to roleplay as a traveler in a medieval fantasy tavern, and you are the wise innkeeper."],
                    ["I want to roleplay as a client meeting with a detective in Victorian London to solve a mystery."],
                    ["I want to roleplay as a student visiting a wise wizard in their magical library."]
                ],
                inputs=roleplay_input
            )
            
            # State variables
            session_state = gr.State({"active": False, "character": None})
            
            def start_roleplay(roleplay_request, current_state):
                """Start new roleplay session"""
                if not roleplay_request.strip():
                    return (
                        current_state,
                        gr.update(),
                        gr.update(),
                        [["System", "Please enter a roleplay scenario first!"]]
                    )
                
                result = self.setup_roleplay(roleplay_request)
                
                if result['status'] == 'success':
                    new_state = {
                        "active": True, 
                        "character": result['character_info']
                    }
                    
                    # Generate opening message
                    opening = self.agent.generate_roleplay_response("*I notice someone new has arrived*")
                    
                    return (
                        new_state,
                        result['character_info'],
                        result['scene_image'],
                        [["🎬 System", result['opening_message']], 
                         ["🎭 " + result['character_info']['character_name'], opening]]
                    )
                else:
                    return (
                        current_state,
                        gr.update(),
                        gr.update(),
                        [["❌ System", f"Error: {result['message']}"]]
                    )
            
            def send_message(message, chat_history, current_state):
                """Send user message and get AI response"""
                if not current_state.get("active", False):
                    chat_history.append(["You", "Please start a roleplay session first!"])
                    return "", chat_history
                
                if message.lower() in ['quit', 'exit', 'end']:
                    response = self.agent.end_roleplay()
                    chat_history.append(["🚪 You", "I need to leave now."])
                    chat_history.append(["🎭 System", response])
                    return "", chat_history
                
                # Generate AI response
                response = self.agent.generate_roleplay_response(message)
                character_name = current_state["character"]["character_name"]
                
                chat_history.append(["💬 You", message])
                chat_history.append([f"🎭 {character_name}", response])
                return "", chat_history
            
            def end_roleplay_session(current_state):
                """End current roleplay session"""
                if current_state.get("active", False):
                    response = self.agent.end_roleplay()
                    new_state = {"active": False, "character": None}
                    return (
                        new_state,
                        gr.update(value=None),
                        gr.update(value=None),
                        [["🎭 System", response]]
                    )
                return current_state, gr.update(), gr.update(), [["🎭 System", "No active session to end."]]
            
            # Event handlers
            setup_btn.click(
                start_roleplay,
                inputs=[roleplay_input, session_state],
                outputs=[session_state, character_info, scene_image, chatbot]
            )
            
            send_btn.click(
                send_message,
                inputs=[user_input, chatbot, session_state],
                outputs=[user_input, chatbot]
            )
            
            user_input.submit(
                send_message,
                inputs=[user_input, chatbot, session_state],
                outputs=[user_input, chatbot]
            )
            
            end_btn.click(
                end_roleplay_session,
                inputs=[session_state],
                outputs=[session_state, character_info, scene_image, chatbot]
            )
        
        return interface

# Example usage
def main():
    # API Keys
    GROQ_API_KEY = "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"  # Replace with your actual key
    
    # Create app (without Stability AI for now)
    app = VisualRoleplayApp(GROQ_API_KEY)
    
    # Launch Gradio interface
    interface = app.create_gradio_interface()
    interface.launch(
        server_name="0.0.0.0",  # Allow external access
        server_port=7860,
        share=True,  # Create public link
        debug=True
    )

if __name__ == "__main__":
    main()