# app.py (New Secure & Robust Version)
import gradio as gr
import httpx
import os
import json

# --- Config ---
# Secrets for our *own* backend
BLAXEL_BASE_URL = os.getenv("BLAXEL_BACKEND_URL")
BLAXEL_API_KEY = os.getenv("BLAXEL_API_KEY") 

# --- Backend Client ---
def call_blaxel_backend(
    user_problem: str, 
    google_key: str, 
    anthropic_key: str, 
    sambanova_key: str
):
    
    # --- 1. Validate Keys ---
    if not BLAXEL_BASE_URL or not BLAXEL_API_KEY:
        yield "Configuration Error: The app's own BLAXEL secrets are not set."
        return
    
    # We require at least the Google key
    if not google_key:
        yield {status_output: "Input Error: Please enter your Google API Key to continue. It is required for the 'Judge' agent."}
        return

    # --- 2. Set up connection details ---
    full_endpoint_url = f"{BLAXEL_BASE_URL.rstrip('/')}/solve_problem"
    
    headers = {
        "Authorization": f"Bearer {BLAXEL_API_KEY}",
        "Content-Type": "application/json"
    }
    
    # --- 3. Securely package the user's keys ---
    payload = {
        "problem": user_problem,
        "keys": {
            "google": google_key,
            "anthropic": anthropic_key or None,  # Send None if empty
            "sambanova": sambanova_key or None, # Send None if empty
        }
    }

    yield f"Connecting to MudabbirAI (at {full_endpoint_url})..."
    
    try:
        # Set a long timeout for all the LLM calls
        with httpx.stream("POST", full_endpoint_url, json=payload, headers=headers, timeout=300.0) as response:
            
            if response.status_code != 200:
                error_details = response.read().decode()
                yield {status_output: f"HTTP Error: Received status code {response.status_code} from Blaxel.\nDetails: {error_details}"}
                return 

            final_json = None
            full_log = ""
            for chunk in response.iter_text():
                if chunk.startswith("FINAL:"):
                    # This is our final package, don't yield it
                    final_json = json.loads(chunk.replace("FINAL:", ""))
                else:
                    # This is a live log update
                    full_log += chunk + "\n"
                    yield {status_output: full_log} # Update log

            # --- 4. Return final results to all UI components ---
            if final_json:
                yield {
                    status_output: full_log + "\nDone!",
                    final_text_output: final_json.get("text"),
                    final_audio_output: final_json.get("audio")
                }

    except httpx.ConnectError as e:
        yield {status_output: f"Connection Error: Could not connect to Blaxel.\nDetails: {str(e)}"}
    except httpx.ReadTimeout:
        yield {status_output: "Error: The connection timed out. (This is a complex agent and may take several minutes.)"}
    except Exception as e:
        yield {status_output: f"An unknown error occurred: {str(e)}"}

# --- Gradio UI (Now with Optional Keys) ---
with gr.Blocks() as demo:
    gr.Markdown("# MudabbirAI: The Self-Calibrating Strategic Selector")
    gr.Markdown("Based on the research by Youssef Hariri (Papers 1 & 2)")

    with gr.Row():
        with gr.Column(scale=2):
            problem_input = gr.Textbox(
                label="Enter Your 'Wicked Problem'",
                placeholder="e.g., 'Develop a new market entry strategy for Southeast Asia...'",
                lines=5
            )
            
            gr.Markdown("### Sponsor API Keys")
            gr.Markdown("Your keys are sent securely, used only for this request, and never stored.")
            
            google_key_input = gr.Textbox(
                label="Google API Key (Required)", type="password",
                placeholder="Required for 'Judge' and 'default' agent"
            )
            
            with gr.Row():
                anthropic_key_input = gr.Textbox(
                    label="Anthropic API Key (Optional)", type="password",
                    placeholder="Enter to enable Anthropic models"
                )
                sambanova_key_input = gr.Textbox(
                    label="SambaNova API Key (Optional)", type="password",
                    placeholder="Enter to enable SambaNova models"
                )
            
            submit_button = gr.Button("Deploy Agent", variant="primary")

        with gr.Column(scale=1):
            status_output = gr.Textbox(
                label="Agent's Internal Monologue (Live Log)",
                lines=20, # Made log even taller
                interactive=False
            )
            final_text_output = gr.Markdown(label="Final Briefing")
            final_audio_output = gr.Audio(label="Audio Presentation")
            
    all_inputs = [
        problem_input,
        google_key_input,
        anthropic_key_input,
        sambanova_key_input
    ]
    
    all_outputs = {
        status_output,
        final_text_output,
        final_audio_output
    }

    submit_button.click(
        fn=call_blaxel_backend,
        inputs=all_inputs,
        outputs=list(all_outputs) # Convert set to list
    )

demo.launch()