import gradio as gr
import requests
import uuid
import os
import json

HF_TOKEN = os.environ.get("HF_TOKEN")

# api server url
API_URL = "https://baxin-simulator.hf.space/simulate"

def generate_unique_name():
    return f"protocol_{uuid.uuid4().hex[:8]}.py"

def simulate_protocol(protocol_content: str):
    if not HF_TOKEN:
        return "Error: HF_TOKEN is not set. Please set it in your environment variables."

    protocol_name = generate_unique_name()
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {HF_TOKEN}"
    }
    
    data = {
        "name": protocol_name,
        "content": protocol_content
    }

    try:
        response = requests.post(API_URL, json=data, headers=headers, timeout=60)
        
        if response.status_code == 200:
            response_data = response.json()
            return f"<pre><code>{response_data.get('run_log', 'Success, but no run log was returned.')}</code></pre>"
        else:
            try:
                error_data = response.json()
                detail = error_data.get("detail", {})
                if isinstance(detail, dict):
                    error_message = detail.get('error_details', response.text)
                else:
                    error_message = detail
                return f"API Error (Status {response.status_code}):\n\n<pre><code>{error_message}</code></pre>"
            except json.JSONDecodeError:
                return f"API Error (Status {response.status_code}):\n\n<pre><code>{response.text}</code></pre>"

    except requests.exceptions.RequestException as e:
        return f"Network Error: Failed to connect to the API.\n\n{e}"

def add_to_chatbot(protocol_text: str, chat_history: list):
    chat_history.append((f"<pre><code>{protocol_text}</code></pre>", None))
    response = simulate_protocol(protocol_text)
    chat_history.append((None, response))
    return chat_history

# --- Gradio UI ---
title = """
<div align="center">
  <h1>Opentrons Protocol Simulator</h1>
  <p>Enter your Opentrons protocol code below and click 'Run Simulation' to see the result on the right.</p>
</div>
"""

with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown(title)
    
    # --- main ---
    with gr.Row(equal_height=False):
        # --- left protocol input ---
        with gr.Column(scale=1):
            textbox = gr.Textbox(
                label="Protocol Input",
                placeholder="from opentrons import protocol_api\n\nmetadata = {'apiLevel': '2.11'}\n\ndef run(protocol: protocol_api.ProtocolContext):\n    tiprack = protocol.load_labware('opentrons_96_tiprack_300ul', 1)\n    p300 = protocol.load_instrument('p300_single_gen2', 'right', tip_racks=[tiprack])\n    p300.pick_up_tip()\n    p300.drop_tip()",
                lines=28,
                show_copy_button=True,
            )
            with gr.Row():
                clear_button = gr.ClearButton([textbox], value="🗑️ Clear Input")
                send_button = gr.Button("▶️ Run Simulation", variant="primary")

        # --- right simulation log ---
        with gr.Column(scale=1):
            chatbot = gr.Chatbot(
                label="Simulation Log", 
                bubble_full_width=True,
                height=650,
                show_copy_button=True,
            )
            gr.ClearButton([chatbot], value="🗑️ Clear Log")


    send_button.click(
        fn=add_to_chatbot,
        inputs=[textbox, chatbot],
        outputs=[chatbot]
    )

app.launch(mcp_server=True)