import gradio as gr
import requests
import os
import uuid
import time

MODAL_BACKEND_URL = os.getenv("MODAL_BACKEND_URL", "").strip()

CUSTOM_CSS = """
:root {
  --primary: #8aa8ff;
  --bg-dark: #1a1a2e;
  --bg-light: #ffffff;
  --text-dark: #e0e0e0;
  --text-light: #2b2b4d;
}
@media (prefers-color-scheme: dark) {
  body, .gradio-container { background-color: var(--bg-dark); color: var(--text-dark); }
  .chatbot-box { background-color: #2b2b4d; border: 1px solid #3a3a5e; border-radius: 12px; padding: 10px; }
  .chatbot-message.bot { background-color: #3a3a5e !important; color: var(--text-dark); padding: 8px; border-radius: 8px; }
  .chatbot-message.user { background: linear-gradient(to right, #4a4a8a, #5c5cb8); color: #fff; padding: 8px; border-radius: 8px; }
}
#send-button {
  background-color: #F97316 !important;
  color: white !important;
  font-weight: 600 !important;
}
#send-button:hover {
  background-color: #EA580C !important;
}
"""

def handle_user_turn(user_input, chat_history, user_role, sid):
    if not user_input.strip():
        return chat_history, ""

    # Gradio's 'messages' type for chatbot expects dictionaries
    # Ensure chat_history is initialized with dictionaries or converted
    if not chat_history or not isinstance(chat_history[0], dict):
        # Convert existing tuples to dictionaries if needed, or start fresh
        new_chat_history = []
        for user_msg, bot_msg in chat_history:
            if user_msg is not None:
                new_chat_history.append({"role": "user", "content": user_msg})
            if bot_msg is not None:
                new_chat_history.append({"role": "assistant", "content": bot_msg})
        chat_history = new_chat_history
    
    chat_history.append({"role": "user", "content": user_input}) # Changed to dictionary format

    if not MODAL_BACKEND_URL:
        # For an error, append as assistant message
        chat_history.append({"role": "assistant", "content": "❗ Configuration error: `MODAL_BACKEND_URL` not set."})
        return chat_history, ""

    thinking_steps = ["Analyzing...", "Searching sources...", "Synthesizing..."]
    
    # Store the last user message to update it later
    # This assumes the last message is always the user's latest input
    last_user_message_index = len(chat_history) - 1

    for step in thinking_steps:
        # Create a temporary 'thinking' message from the assistant
        temp_message = {"role": "assistant", "content": f"<div class='chatbot-message bot'><em>{step}</em></div>"}
        
        # If there's an existing assistant thinking message, update it in place
        # Otherwise, append a new one
        if len(chat_history) > last_user_message_index + 1 and chat_history[last_user_message_index + 1]["role"] == "assistant":
            chat_history[last_user_message_index + 1] = temp_message
        else:
            chat_history.append(temp_message)

        yield chat_history, "" # Yield with the temporary thinking message

        time.sleep(0.4)

    try:
        response = requests.post(MODAL_BACKEND_URL, json={
            "user_input": user_input,
            "session_id": sid,
            "role": user_role
        }, timeout=600)
        response.raise_for_status()
        reply = response.json().get("final_markdown", "No response received.")
    except Exception as e:
        reply = f"❗ Error: {str(e)}"
    
    # Update the last assistant message with the final reply
    # This should replace the thinking message or be appended if no thinking message was added
    if len(chat_history) > last_user_message_index + 1 and chat_history[last_user_message_index + 1]["role"] == "assistant":
        chat_history[last_user_message_index + 1] = {"role": "assistant", "content": reply}
    else:
        chat_history.append({"role": "assistant", "content": reply}) # Should not typically happen if thinking steps run

    yield chat_history, ""

with gr.Blocks(css=CUSTOM_CSS, title="MediAgent AI") as demo:
    session_id = gr.State(lambda: str(uuid.uuid4()))

    gr.Markdown("""
    # 🧬 MediAgent AI
    **Your AI Health Navigator**

    <div class='chatbot-box'>
    MediAgent AI is a multilingual, role-based assistant that gives clear, cited medical answers.
    <ul>
      <li>🗣 Ask in English, Spanish, Hindi, or French</li>
      <li>🎯 <b>Patient</b> = simplified advice | <b>Doctor</b> = in-depth, technical insights</li>
      <li>🔗 Answers are backed by PubMed, FDA, etc.</li>
      <li>💊 Can provide information on **drug side effects, interactions, and official label details.**</li>
      <li>🔬 Can search for **clinical trials and medical research.**</li>
      <li>🛡️ Detects emergencies, adds safety notices</li>
    </ul>
    </div>
    """) # Changed Markdown content

    with gr.Row():
        with gr.Column(scale=1):
            user_role = gr.Radio(
                choices=["patient", "doctor"],
                value="patient",
                label="Select Role",
                info="Choose your user type for tailored responses"
            )
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(label="Conversation", height=480, type='messages') # Added type='messages'

    with gr.Row():
        user_textbox = gr.Textbox(
            placeholder="Ask your health question...",
            lines=1,
            scale=4
        )
        submit_btn = gr.Button("Send", elem_id="send-button", scale=1)

    examples = gr.Examples(
        examples=[
            "What are the side effects of metformin?",
            "Explain Crohn’s disease vs Ulcerative Colitis for a doctor",
            "¿Cuáles son los síntomas de la hipertensión?",
            "Can I take ibuprofen with lisinopril? (for interactions)", # New example
            "What is the FDA label information for acetaminophen?", # New example
            "Are there any ongoing clinical trials for breast cancer?", # New example
        ],
        inputs=[user_textbox]
    ) # Changed Examples content

    gr.ClearButton([chatbot, user_textbox], value="🔄 Clear")

    submit_btn.click(
        fn=handle_user_turn,
        inputs=[user_textbox, chatbot, user_role, session_id],
        outputs=[chatbot, user_textbox]
    )

    user_textbox.submit(
        fn=handle_user_turn,
        inputs=[user_textbox, chatbot, user_role, session_id],
        outputs=[chatbot, user_textbox]
    )

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