import gradio as gr
import os
import requests
from datetime import datetime

# Mistral API setup
MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY", "your_api_key_here")
MISTRAL_API_URL = "https://api.mistral.ai/v1/chat/completions"

# Local medical knowledge base
MEDICAL_KNOWLEDGE = {
    "headache": {
        "assessment": "Could be tension-type, migraine, or other common headache",
        "recommendations": [
            "Rest in a quiet, dark room",
            "Apply cold compress to forehead",
            "Hydrate well (at least 8 glasses water/day)",
            "Try gentle neck stretches",
            "Consider OTC pain relievers (ibuprofen 200-400mg or paracetamol 500mg)"
        ],
        "warnings": [
            "If sudden/severe 'thunderclap' headache",
            "If with fever, stiff neck, or confusion",
            "If following head injury"
        ]
    },
    "fever": {
        "assessment": "Common symptom of infection or inflammation",
        "recommendations": [
            "Rest and increase fluid intake",
            "Use light clothing and keep room cool",
            "Sponge with lukewarm water if uncomfortable",
            "Paracetamol 500mg every 6 hours (max 4 doses/day)",
            "Monitor temperature every 4 hours"
        ],
        "warnings": [
            "Fever >39°C lasting >3 days",
            "If with rash, difficulty breathing, or seizures",
            "In infants <3 months with any fever"
        ]
    },
    "sore throat": {
        "assessment": "Often viral, but could be strep throat",
        "recommendations": [
            "Gargle with warm salt water 3-4 times daily",
            "Drink warm liquids (tea with honey)",
            "Use throat lozenges",
            "Stay hydrated",
            "Rest your voice"
        ],
        "warnings": [
            "Difficulty swallowing or breathing",
            "White patches on tonsils",
            "Fever >38.5°C with throat pain"
        ]
    }
}

def get_local_advice(symptom):
    """Get structured medical advice from local knowledge base"""
    symptom_lower = symptom.lower()
    for key in MEDICAL_KNOWLEDGE:
        if key in symptom_lower:
            knowledge = MEDICAL_KNOWLEDGE[key]
            response = "🔍 [Assessment]\n- " + knowledge["assessment"] + "\n\n"
            response += "💡 [Self-Care Recommendations]\n" + "\n".join(["• " + item for item in knowledge["recommendations"]]) + "\n\n"
            response += "⚠️ [When to Seek Medical Care]\n" + "\n".join(["• " + item for item in knowledge["warnings"]]) + "\n\n"
            response += "📅 [Follow-up]\n• Re-evaluate in 2-3 days if not improving\n• See doctor if symptoms worsen or persist beyond 5 days"
            return response
    return None

def query_mistral(prompt):
    """Try API first, fallback to local knowledge"""
    try:
        headers = {"Authorization": f"Bearer {MISTRAL_API_KEY}", "Content-Type": "application/json"}
        payload = {
            "model": "mistral-tiny",
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.7,
            "max_tokens": 1024
        }
        response = requests.post(MISTRAL_API_URL, headers=headers, json=payload, timeout=10)
        response.raise_for_status()
        return response.json()['choices'][0]['message']['content']
    except:
        return None

def generate_response(name, message):
    """Generate medical response with fallback"""
    # First try local knowledge base
    local_response = get_local_advice(message)
    if local_response:
        return f"Hello {name},\n\n{local_response}\n\nWishing you good health,\nDr. Alex"
    
    # Then try API
    prompt = f"""As Dr. Alex, provide structured advice for {name} who reports: "{message}"

    Format response with these sections:
    
    [Assessment] - Brief professional evaluation
    [Self-Care Recommendations] - 3-5 specific actionable steps
    [When to Seek Medical Care] - Clear warning signs
    [Follow-up] - Monitoring advice
    
    Use bullet points, professional yet compassionate tone."""
    
    api_response = query_mistral(prompt)
    if api_response:
        return api_response
    
    # Final fallback
    common_conditions = "\n".join([f"- {cond}" for cond in MEDICAL_KNOWLEDGE.keys()])
    return f"""Hello {name},

I'm currently unable to access detailed medical databases. For general advice:

Common conditions I can advise on:
{common_conditions}

For immediate concerns:
• Contact your local healthcare provider
• In emergencies, call your local emergency number

Please try asking about one of these common conditions or consult a healthcare professional.

Best regards,
Dr. Alex"""

# Gradio Interface
with gr.Blocks(title="Dr. Alex Medical Advisor", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🩺 Dr. Alex - General Health Advisor")
    name = gr.Textbox(label="Your Name", placeholder="Enter your name")
    
    with gr.Tab("Health Consultation"):
        gr.Markdown("## Describe your health concern")
        msg = gr.Textbox(label="Symptoms/Question", lines=3)
        submit_btn = gr.Button("Get Medical Advice")
        output = gr.Textbox(label="Dr. Alex's Advice", lines=15, interactive=False)
    
    submit_btn.click(
        fn=lambda name, msg: generate_response(name, msg),
        inputs=[name, msg],
        outputs=output
    )
    
    gr.Markdown("### Common Health Topics")
    with gr.Row():
        for condition in MEDICAL_KNOWLEDGE:
            gr.Button(condition.capitalize()).click(
                fn=lambda name, cond=condition: generate_response(name, cond),
                inputs=[name],
                outputs=output
            )

demo.launch()