import os
import gradio as gr
from anthropic import Anthropic
import requests
from dotenv import load_dotenv
import time

# Load environment variables
load_dotenv()

# Initialize Anthropic client
client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))

MODAL_CLINIC_ENDPOINT = "https://aayushraj0324--healthmate-clinic-lookup-search-clinics.modal.run"

def classify_urgency(symptoms: str) -> str:
    """Classify the urgency level of the symptoms using Claude."""
    prompt = f"""You are a medical triage assistant. Given this symptom description: {symptoms}, \nclassify it as: emergency / routine visit / home care. Explain briefly."""
    message = client.messages.create(
        model="claude-sonnet-4-20250514",
        #model="claude-3-sonnet-20240229",
        max_tokens=200,
        temperature=0.01,
        system="You are a medical triage assistant. Provide clear, concise classifications.",
        messages=[{"role": "user", "content": prompt}]
    )
    return message.content[0].text

def get_possible_conditions(symptoms: str) -> str:
    """Get possible medical conditions based on symptoms using Claude."""
    prompt = f"""List 2–4 possible medical conditions that match these symptoms: {symptoms}. \nKeep it non-technical and easy to understand."""
    message = client.messages.create(
        model="claude-sonnet-4-20250514",
        #model="claude-3-sonnet-20240229",
        max_tokens=400,
        temperature=0.01,
        system="You are a medical assistant. Provide clear, non-technical explanations of possible conditions.",
        messages=[{"role": "user", "content": prompt}]
    )
    return message.content[0].text

def lookup_clinics(city: str) -> str:
    return "Clinic Lookup will be back soon"
    # try:
    #     response = requests.get(MODAL_CLINIC_ENDPOINT, params={"city": city}, timeout=20)
    #     response.raise_for_status()
    #     clinics = response.json()
    #     if clinics and isinstance(clinics, list) and "error" not in clinics[0]:
    #         return "\n\n".join([
    #             f"🏥 {clinic['name']}\n🔗 {clinic['link']}\n📝 {clinic['description']}"
    #             for clinic in clinics
    #         ])
    #     else:
    #         return clinics[0].get("error", "No clinics found.")
    # except Exception as e:
    #     return f"Error finding clinics: {str(e)}"

def process_input(symptoms: str, city: str, pain_level: int, life_impact: int) -> tuple:
    """Process the input and return all results."""
    time.sleep(1)

    # Enrich symptom input with slider info
    enriched_input = (
        f"{symptoms}\n"
        f"Pain level: {pain_level}/10\n"
        f"Impact on daily life: {life_impact}/10"
    )

    # Use enriched input for Claude
    urgency = classify_urgency(enriched_input)
    conditions = get_possible_conditions(enriched_input)

    # Nearby clinics
    if city:
        clinic_text = lookup_clinics(city)
    else:
        clinic_text = "Please provide a city to find nearby clinics."
    
    urgency_md = f"### 🩺 Urgency Classification\n\n{urgency}\n\n---"
    conditions_md = f"### ❓ Possible Conditions\n\n{conditions}\n\n---"
    clinics_md = f"### 🏥 Nearby Clinics\n\n{clinic_text}"
    return urgency_md, conditions_md, clinics_md

def full_handler(symptoms, city, pain_level, life_impact):
    status_text = "⏳ Processing..."
    urgency_md, conditions_md, clinics_md = process_input(symptoms, city, pain_level, life_impact)
    return status_text, urgency_md, conditions_md, clinics_md, ""  # Last "" clears status

# Create the Gradio interface
with gr.Blocks(css=".gradio-container {max-width: 800px; margin: auto;}") as demo:
    gr.Markdown(
        """
        # 🏥 HealthMate: AI Medical Triage Assistant
        Enter your symptoms and optionally your city to get medical guidance and nearby clinic recommendations.
        """
    )
    
    with gr.Row():
        with gr.Column():
            symptoms = gr.Textbox(
                label="Describe your symptoms",
                placeholder="Example: I have a severe headache and fever for the past 2 days...",
                lines=4
            )
            pain_level = gr.Slider(
                minimum=0, maximum=10, step=1, value=5,
                label="Pain Level (0 = none, 10 = unbearable)"
            )
            life_impact = gr.Slider(
                minimum=0, maximum=10, step=1, value=5,
                label="Impact on Daily Life (0 = no impact, 10 = can't function)"
            )
            city = gr.Textbox(
                label="Your city (optional)",
                placeholder="Example: San Francisco"
            )
            submit_btn = gr.Button("Get Medical Guidance", variant="primary")
            status = gr.Markdown()
    
    with gr.Row():
        with gr.Column() as outputs:
            urgency = gr.Markdown()
            conditions = gr.Markdown()
            clinics = gr.Markdown()
    
    submit_btn.click(
        fn=full_handler,
        inputs=[symptoms, city, pain_level, life_impact],
        outputs=[status, urgency, conditions, clinics, status],
        show_progress=True
    )

    footerMD = """⚠️ Disclaimer: HealthMate is an AI-based assistant
             and not a substitute for professional medical advice,
             diagnosis, or treatment. Always seek the advice of a 
             qualified healthcare provider with any questions you 
             may have regarding a medical condition. If you think 
             you may have a medical emergency, call your doctor or 
             local emergency services immediately.
    """
    gr.Markdown(footerMD, elem_classes="footer")

if __name__ == "__main__":
    demo.launch(share=True, pwa=True)