import gradio as gr
import requests
import os
import uuid

# Configuration
LANGFLOW_API_URL = os.environ.get("LANGFLOW_API_URL", "")
LANGFLOW_API_KEY = os.environ.get("LANGFLOW_API_KEY", "")
HF_API_KEY = os.environ.get("HF_API_KEY", "")

# Dictionary to store session IDs per Gradio session
# Key: Gradio session hash, Value: Langflow session ID
session_storage = {}

def call_langflow(message, history, request: gr.Request = None):
    """
    Call Langflow API and return the response with session persistence
    """
    # Get or create session ID for this user
    # Use a default session if request is None (happens during example caching)
    if request is None:
        session_hash = "default"
    else:
        session_hash = request.session_hash
    
    if session_hash not in session_storage:
        session_storage[session_hash] = str(uuid.uuid4())
        print(f"🆕 NEW SESSION created: {session_storage[session_hash]}")
    
    langflow_session_id = session_storage[session_hash]

    # Debug logging
    print(f"📤 Sending message: {message[:50]}...")  # First 50 chars
    print(f"🔑 Using session ID: {langflow_session_id}")
    print(f"👤 Gradio session hash: {session_hash}")
    
    headers = {
        "Content-Type": "application/json",
    }
    
    # Add API keys
    if HF_API_KEY:
        headers["Authorization"] = f"Bearer {HF_API_KEY}"
    if LANGFLOW_API_KEY:
        headers["x-api-key"] = f"{LANGFLOW_API_KEY}"
    
    # Adjust this payload based on your Langflow API structure
    payload = {
        "input_value": message,
        "output_type": "chat",
        "input_type": "chat",
        "session_id": langflow_session_id,  # Add session ID
        "tweaks": {}
    }
    
    try:
        response = requests.post(
            LANGFLOW_API_URL,
            json=payload,
            headers=headers,
            timeout=30
        )
        response.raise_for_status()
        
        # Parse response - adjust based on your API response structure
        data = response.json()
        
        # Common Langflow response structures:
        # Option 1: data["outputs"][0]["outputs"][0]["results"]["message"]["text"]
        # Option 2: data["result"]["message"]
        # Adjust the following line based on your actual response:
        
        bot_message = data["outputs"][0]["outputs"][0]["results"]["message"]["text"]
        return bot_message
        
    except requests.exceptions.RequestException as e:
        return f"Error connecting to Langflow: {str(e)}"
    except (KeyError, IndexError) as e:
        return f"Error parsing response: {str(e)}\nResponse: {data}"

# Create Gradio Chat Interface
custom_theme = gr.themes.Default(
    primary_hue="pink",       # main primary color (affects submit buttons)
    #secondary_hue="blue",    # secondary highlights (hover, accents)
    font="Arial",             # optional font
)

demo = gr.ChatInterface(
    fn=call_langflow,
    title="Urban Air Chatbot POC",
    description="Ask Urbie about Urban Air, Westminster.",
    examples=["What attractions do you have?", "Can you give me details about your membership plans?"],
    theme=custom_theme,
    retry_btn=None,
    undo_btn=None,  # Keep it as None or remove the line
    cache_examples=False,  # Disable caching to avoid startup issues
    css="""
        h1 { color: yellow !important; }
    """)

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