# In app.py

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: str, chat_history: list, user_role: str, sid: str):
    """
    Handles a single turn of the conversation.
    """
    if not user_input.strip():
        return chat_history, ""

    # Initialize chat_history if it's the first turn (it will be None)
    if chat_history is None:
        chat_history = []
    
    # Append the user's message in the correct dictionary format
    chat_history.append({"role": "user", "content": user_input})
    
    # Immediately yield the updated history to show the user's message
    yield chat_history, ""

    if not MODAL_BACKEND_URL:
        chat_history.append({"role": "assistant", "content": "❗ Configuration error: `MODAL_BACKEND_URL` is not set."})
        yield chat_history, ""
        return

    # --- Streaming "Thinking" Animation ---
    chat_history.append({"role": "assistant", "content": ""}) # Add an empty placeholder for the bot response
    
    thinking_steps = ["Analyzing your query...", "Searching medical databases...", "Consulting research papers...", "Synthesizing the answer..."]
    for step in thinking_steps:
        # Update the last message (the bot's placeholder) with the current thinking step
        chat_history[-1]["content"] = f"<em>{step}</em>"
        yield chat_history, ""
        time.sleep(0.4)
    # --- End of Animation ---

    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 valid response was received from the backend.")
    except requests.exceptions.RequestException as e:
        reply = f"❗ **Error Communicating with Backend:**\n\nCould not get a response. The server may be starting up or experiencing an issue. Please try again in a moment. \n\n*Details: {str(e)}*"
    except Exception as e:
        reply = f"❗ **An Unexpected Error Occurred:**\n\n*Details: {str(e)}*"
    
    # Update the final bot message with the actual reply
    chat_history[-1]["content"] = reply

    yield chat_history, ""


# --- The rest of your Gradio app code remains the same ---
with gr.Blocks(css=CUSTOM_CSS, title="MediAgent AI: Conversational Health Navigator & MCP Tool Server") as demo:
    session_id = gr.State(lambda: str(uuid.uuid4()))

    gr.Markdown("""
    # 🧬 MediAgent AI: Conversational Health Navigator & MCP Tool Server
    
    <div class='chatbot-box'>
    <b>MediAgent AI</b> is a multilingual, role-adaptive health assistant and a fully MCP-compliant tool server.<br>
    <ul>
      <li>🌐 <b>Multilingual:</b> English, Spanish, Hindi, French</li>
      <li>👤 <b>Role-Based:</b> <b>Patient</b> = simplified advice | <b>Doctor</b> = in-depth, technical insights</li>
      <li>🛡️ <b>Critical Safety Layer:</b> Detects emergencies, adds safety notices</li>
      <li>🔗 <b>Sourced & Verifiable:</b> Answers are backed by PubMed, FDA, etc.</li>
      <li>💬 <b>Stateful Conversations:</b> Remembers chat history for contextual follow-up</li>
      <li>🛠️ <b>MCP Tool Server:</b> All tools are externally callable via MCP</li>
      <li>💊 <b>Drug Info:</b> Side effects, interactions, FDA label details</li>
      <li>🔬 <b>Research:</b> Clinical trials & medical literature search</li>
    </ul>
    <div style='margin-top:10px; text-align:center; color:#8aa8ff; font-weight:bold;'>Powered by Modal and LlamaIndex. Special thanks for their support!</div>
    </div>
    """)

    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):
            # IMPORTANT: Ensure chatbot is defined with type='messages'
            chatbot = gr.Chatbot(label="Conversation", height=480, 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, variant="primary")

    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)",
            "What is the FDA label information for acetaminophen?",
            "Are there any ongoing clinical trials for breast cancer?",
        ],
        inputs=[user_textbox]
    )

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

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


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