import gradio as gr
import os
import httpx
import json
import base64
import uuid
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Configuration
AGENT_URL = os.getenv("BLAXEL_AGENT_URL", "https://run.blaxel.ai/silvestre-po/agents/curiosity-agent/generate")
BLAXEL_API_KEY = os.getenv("BLAXEL_API_KEY")

# Frontend Rate Limiting (configurable via env var)
import time
from collections import deque
# Default: 1 request per minute (60 seconds)
# For overnight/low-traffic: set RATE_LIMIT_WINDOW=3600 (1 hour)
FRONTEND_RATE_LIMIT_WINDOW = int(os.getenv("RATE_LIMIT_WINDOW", "60"))
FRONTEND_RATE_LIMIT_MAX = 1
frontend_request_history = deque(maxlen=FRONTEND_RATE_LIMIT_MAX)

def check_frontend_rate_limit():
    """Returns True if rate limit would be exceeded, False otherwise."""
    now = time.time()
    # Remove old requests outside the window
    while frontend_request_history and (now - frontend_request_history[0]) > FRONTEND_RATE_LIMIT_WINDOW:
        frontend_request_history.popleft()
    
    # Check if we're at the limit
    if len(frontend_request_history) >= FRONTEND_RATE_LIMIT_MAX:
        time_to_wait = FRONTEND_RATE_LIMIT_WINDOW - (now - frontend_request_history[0])
        return True, time_to_wait
    
    # Record this request
    frontend_request_history.append(now)
    return False, 0

async def generate_storybook(question, age, session_id, progress=gr.Progress()):
    """
    Calls the Blaxel Agent to generate the storybook page using Streaming Response.
    Yields updates to the UI as data arrives.
    Context Engineering: Maintains session and displays related questions.
    """
    # Frontend rate limiting check
    is_limited, wait_time = check_frontend_rate_limit()
    if is_limited:
        error_msg = f"⏳ Please wait {int(wait_time)} seconds. The app is limited to 1 story per minute to manage costs."
        yield None, None, None, error_msg, gr.update(choices=[], value=None)
        return
    
    if not question:
        yield None, None, None, "Please ask a question!", gr.update(choices=[], value=None)
        return

    print(f"🚀 Sending request for: '{question}' (Age: {age}) [Session: {session_id[:8]}...]")
    
    headers = {
        "Content-Type": "application/json"
    }
    
    # Security: Shared Secret (Bearer Token)
    shared_secret = os.getenv("CURIOSITY_SHARED_SECRET")
    if shared_secret:
        headers["Authorization"] = f"Bearer {shared_secret}"
    elif BLAXEL_API_KEY:
        # Fallback to Blaxel key if no specific secret set (though they are different concepts)
        # Ideally we just use the shared secret for the agent auth
        pass

    payload = {
        "question": question,
        "age": int(age),
        "session_id": session_id  # Context Engineering
    }

    # Initialize state
    current_story = ""
    current_image = None
    current_audio = None
    related_questions = []  # Context Engineering
    status_msg = "🚀 **Starting your magical story...**"
    
    # Progress tracking
    progress(0, desc="Starting...")
    
    try:
        async with httpx.AsyncClient(timeout=600.0) as client:
            async with client.stream("POST", AGENT_URL, json=payload, headers=headers) as response:
                if response.status_code != 200:
                    yield None, None, None, f"Error: Agent returned {response.status_code}", gr.update(choices=[], value=None)
                    return

                step = 0
                total_steps = 3  # Story, Image, Audio
                
                async for line in response.aiter_lines():
                    if not line:
                        continue
                        
                    try:
                        chunk = json.loads(line)
                        
                        # 1. Status Update
                        if "status" in chunk:
                            status_msg = f"🔄 {chunk['status']}"
                            if "heartbeat" in chunk['status']:
                                progress((step / total_steps) + (0.33 * 0.5), desc=status_msg)
                            yield current_story, current_image, current_audio, status_msg, gr.update(choices=related_questions if related_questions else [], value=None)
                        
                        # 2. Story Data
                        if "data" in chunk:
                            data = chunk["data"]
                            if isinstance(data, list) and len(data) > 0:
                                data = data[0]
                                
                            if isinstance(data, dict):
                                title = data.get("title", "Untitled")
                                story = data.get("story", "")
                                current_story = f"# {title}\n\n{story}"
                                status_msg = "📖 Story generated!"
                                step = 1
                                progress(step / total_steps, desc=status_msg)
                                yield current_story, current_image, current_audio, status_msg, gr.update(choices=related_questions if related_questions else [], value=None)
                        
                        # 3. Image Data
                        if "image_data" in chunk:
                            try:
                                b64_data = chunk["image_data"]
                                if "," in b64_data:
                                    b64_data = b64_data.split(",")[1]
                                    
                                image_bytes = base64.b64decode(b64_data)
                                import io
                                from PIL import Image
                                current_image = Image.open(io.BytesIO(image_bytes))
                                status_msg = "🎨 Image generated!"
                                step = 2
                                progress(step / total_steps, desc=status_msg)
                                yield current_story, current_image, current_audio, status_msg, gr.update(choices=related_questions if related_questions else [], value=None)
                            except Exception as e:
                                print(f"Image decode error: {e}")
                        
                        # 4. Audio Data
                        if "audio_data" in chunk:
                            try:
                                b64_data = chunk["audio_data"]
                                if "," in b64_data:
                                    b64_data = b64_data.split(",")[1]
                                    
                                audio_bytes = base64.b64decode(b64_data)
                                import tempfile
                                with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as fp:
                                    fp.write(audio_bytes)
                                    current_audio = fp.name
                                status_msg = "🔊 Audio generated!"
                                step = 3
                                progress(1.0, desc=status_msg)
                                yield current_story, current_image, current_audio, status_msg, gr.update(choices=related_questions if related_questions else [], value=None)
                            except Exception as e:
                                print(f"Audio decode error: {e}")
                        
                        # 5. Related Questions (Context Engineering)
                        if "related_questions" in chunk:
                            related_questions = chunk.get("related_questions", [])
                            is_followup = chunk.get("is_followup", False)
                            if is_followup:
                                status_msg = "🔗 Follow-up question detected! Generating context-aware suggestions..."
                            else:
                                status_msg = "💡 Generating related questions..."
                            # Update Radio with new choices
                            yield current_story, current_image, current_audio, status_msg, gr.update(choices=related_questions, value=None)
                        
                        # 6. Errors
                        if "error" in chunk:
                            status_msg = f"❌ Oops! Something went wrong. Please try again!"
                            yield current_story, current_image, current_audio, status_msg, gr.update(choices=related_questions if related_questions else [], value=None)
                            
                    except json.JSONDecodeError:
                        pass
                        
        progress(1.0, desc="✅ Done!")
        final_status = "✅ All done! Enjoy your story! 🎉"
        if related_questions:
            final_status += "\n\n💡 Click on a related question below to continue exploring!"
        yield current_story, current_image, current_audio, final_status, gr.update(choices=related_questions, value=None)

    except Exception as e:
        yield current_story, current_image, current_audio, f"😢 Connection Error: Please check your internet and try again!", gr.update(choices=[], value=None)

# Enhanced Custom CSS for kids
custom_css = """
.container { max-width: 1000px; margin: auto; }
h1 { 
    text-align: center; 
    color: #FF6B6B; 
    font-family: 'Comic Sans MS', 'Chalkboard SE', 'Arial Rounded MT Bold', sans-serif;
    font-size: 3em !important;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
}
h3 {
    font-family: 'Comic Sans MS', 'Chalkboard SE', sans-serif;
    font-size: 1.3em !important;
    color: #555;
}
.gradio-button-primary { 
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: bold !important;
    font-size: 1.3em !important;
    padding: 15px 30px !important;
    border-radius: 15px !important;
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.4) !important;
    transition: all 0.3s ease !important;
}
.gradio-button-primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(102, 126, 234, 0.6) !important;
}
.gradio-button-primary:disabled {
    background: linear-gradient(135deg, #ccc 0%, #999 100%) !important;
    cursor: not-allowed !important;
    transform: none !important;
}
label {
    font-family: 'Comic Sans MS', 'Chalkboard SE', sans-serif !important;
    font-size: 1.1em !important;
    color: #333 !important;
    font-weight: bold !important;
}
"""

''' with gr.Blocks(title="🦁 Curiosity Storybook") as demo:
    gr.Markdown("""
# 🦁 Curiosity Storybook

[![Blaxel](https://img.shields.io/badge/Agent-Blaxel-purple)](https://blaxel.ai) [![Blaxel](https://img.shields.io/badge/MCP-Blaxel-purple)](https://blaxel.ai) [![Modal](https://img.shields.io/badge/Compute-Modal-green)](https://modal.com) [![Gemini](https://img.shields.io/badge/LLM-Gemini%202.5%20Pro-blue)](https://deepmind.google/technologies/gemini/) [![Hyperbolic](https://img.shields.io/badge/LLM-Llama%203.3-orange)](https://hyperbolic.xyz) [![Flux](https://img.shields.io/badge/Image-Flux.1-red)](https://blackforestlabs.ai/) [![OpenAI](https://img.shields.io/badge/Audio-OpenAI%20TTS-green)](https://openai.com) [![Gradio](https://img.shields.io/badge/Frontend-Gradio-orange)](https://gradio.app)

### Ask me anything! I'll remember our conversation and suggest related questions. 📚✨💭
*Works in any language! Each conversation builds on previous questions.*
""")  '''

# --- 1. Definición del Contenido y Estilos CSS ---

# A. CSS para Adaptabilidad: Permite el salto de línea (wrap)
COMPRESSION_CSS_ADAPTABLE = """
<style>
/* Clase para permitir que los elementos se envuelvan (wrap) naturalmente */
.compact-badges-adaptable {
    padding: 0;
    margin: 0;
}

/* Estilo para los enlaces (los badges) */
.compact-badges-adaptable a {
    /* Es crucial que sean inline-block para que se alineen horizontalmente */
    display: inline-block !important; 
    /* Margen entre badges: 4px horizontal, 5px vertical (para separar las filas) */
    margin: 0px 4px 5px 0px !important; 
    padding: 0;
    vertical-align: middle;
}
</style>
"""

# B. Badges Convertidos a HTML Puro para garantizar la visualización
# Nota: La sintaxis es <a href="URL_DESTINO"><img src="URL_IMAGEN" alt="Texto" style="height: 20px;"></a>
BADGES_HTML_PURO = """
<a href="https://blaxel.ai"><img src="https://img.shields.io/badge/Agent-Blaxel-purple" alt="Blaxel Agent" style="height: 20px;"/></a>
<a href="https://blaxel.ai"><img src="https://img.shields.io/badge/MCP-Blaxel-purple" alt="Blaxel MCP" style="height: 20px;"/></a>
<a href="https://modal.com"><img src="https://img.shields.io/badge/Compute-Modal-green" alt="Modal Compute" style="height: 20px;"/></a>
<a href="https://deepmind.google/technologies/gemini/"><img src="https://img.shields.io/badge/LLM-Gemini%202.5%20Pro-blue" alt="Gemini LLM" style="height: 20px;"/></a>
<a href="https://hyperbolic.xyz"><img src="https://img.shields.io/badge/LLM-Llama%203.3-orange" alt="Llama LLM" style="height: 20px;"/></a>
<a href="https://blackforestlabs.ai/"><img src="https://img.shields.io/badge/Image-Flux.1-red" alt="Flux Image" style="height: 20px;"/></a>
<a href="https://openai.com"><img src="https://img.shields.io/badge/Audio-OpenAI%20TTS-green" alt="OpenAI Audio" style="height: 20px;"/></a>
<a href="https://gradio.app"><img src="https://img.shields.io/badge/Frontend-Gradio-orange" alt="Gradio Frontend" style="height: 20px;"/></a>
"""

# C. Ensamblamos el contenido final que inyectaremos en gr.HTML
FULL_HTML_ADAPTABLE = f"""
{COMPRESSION_CSS_ADAPTABLE}
<div class="compact-badges-adaptable">
    {BADGES_HTML_PURO} 
</div>
"""

# --- 2. Estructura Principal de Gradio ---

with gr.Blocks(title="🦁 Curiosity Storybook") as demo:
    gr.Markdown("""
# 🦁 Curiosity Storybook
""")
    
    # Inyectamos el CSS y el contenido de los badges
    gr.HTML(FULL_HTML_ADAPTABLE)

    gr.Markdown("""
---
### Ask me anything! I'll remember our conversation and suggest related questions. 📚✨💭
*Works in any language! Each conversation builds on previous questions.*
""")

    




    # Session state (invisible to user, persistent across interactions)
    session_state = gr.State(lambda: str(uuid.uuid4()))
    
    with gr.Row():
        with gr.Column(scale=1):
            question_input = gr.Textbox(
                label="❓ What are you curious about?", 
                placeholder="e.g., Why is the sky blue? 🌤️\n¿Por qué brillan las estrellas? ⭐\nPourquoi les chats ronronnent? 🐱",
                lines=3,
                max_lines=5
            )
            age_slider = gr.Slider(
                minimum=3, 
                maximum=12, 
                value=7, 
                step=1, 
                label="🎂 Child's Age"
            )
            
            # Context Engineering: Related questions UI
            related_questions_box = gr.Radio(
                label="💡 Related Questions (click to ask)", 
                choices=[],
                interactive=True,
                visible=True
            )
            
            submit_btn = gr.Button("✨ Tell me a story! ✨", variant="primary", size="lg")
            status_output = gr.Markdown("🌟 **Ready to explore!** Ask me anything...")
            
        with gr.Column(scale=2):
            story_output = gr.Markdown(label="📖 Your Story")
            with gr.Row():
                image_output = gr.Image(label="🎨 Illustration", type="pil")
                audio_output = gr.Audio(
                    label="🔊 Narration", 
                    type="filepath",
                    autoplay=True
                )
    
    # Pre-loaded examples
    gr.Examples(
        examples=[
            ["Why is the sky blue?", 7],
            ["¿Por qué brillan las estrellas?", 8],
            ["Pourquoi les chats ronronnent?", 6],
            ["Why do we dream?", 9],
            ["¿Cómo vuelan los aviones?", 5],
            ["Why does the moon change shape?", 7],
        ],
        inputs=[question_input, age_slider],
        label="💡 Try these questions:"
    )
    
    # Footer with credits
    gr.Markdown("""
---
### 🏆 Built for MCP's 1st Birthday Hackathon


*Featuring Context Engineering with Compaction & Relevance Detection*
""")

    # Event: Submit button click
    submit_btn.click(
        fn=generate_storybook,
        inputs=[question_input, age_slider, session_state],
        outputs=[story_output, image_output, audio_output, status_output, related_questions_box],
        show_progress="full",
        concurrency_limit=1
    )
    
    # Event: Click on related question
    def select_related_question(selected):
        """When user clicks a related question, auto-fill the input."""
        if selected:
            return selected
        return ""
    
    related_questions_box.change(
        fn=select_related_question,
        inputs=[related_questions_box],
        outputs=[question_input]
    ).then(
        fn=generate_storybook,
        inputs=[question_input, age_slider, session_state],
        outputs=[story_output, image_output, audio_output, status_output, related_questions_box],
        show_progress="full",
        concurrency_limit=1
    )

if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(),
        css=custom_css  # CSS va aquí en Gradio 6
    )