import gradio as gr
import google.generativeai as genai
import os
from dotenv import load_dotenv

# ✅ Load environment variables (for local .env use)
load_dotenv()

# ✅ Fetch Gemini API Key (works for both local + Hugging Face + Render)
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

if not GEMINI_API_KEY:
    raise ValueError("❌ GEMINI_API_KEY not found. Please set it in your .env or Hugging Face secrets.")

# 🔹 Configure Gemini client
genai.configure(api_key=GEMINI_API_KEY)
MODEL_ID = "gemini-2.0-flash"  # or gemini-1.5-flash for faster inference

# ---------------- AI Response Function ----------------
def respond(albumin, creatinine, glucose, crp, mcv, rdw, alp, wbc, lymphocytes,
            hemoglobin, pv, age, gender, height, weight):

    system_message = (
        "You are an AI Health Assistant that analyzes laboratory biomarkers "
        "and generates beautifully formatted, easy-to-read health summaries.\n\n"
    
        "FORMATTING RULES - EXTREMELY IMPORTANT:\n"
        "• Use clear section headers with emojis\n"
        "• Add blank lines between sections for better readability\n"
        "• Use bullet points (•) instead of dashes\n"
        "• Use bold (**text**) for emphasis on important values\n"
        "• Use status emojis: 🟢 Normal | 🟡 Monitor | 🔴 Needs Attention\n"
        "• Keep paragraphs short (2-3 sentences max)\n"
        "• Add visual separators (---) between major sections\n\n"
            
        "REQUIRED STRUCTURE:\n\n"
        "# 📊 Your Health Report\n\n"
        "---\n\n"
        "## 🔬 Biomarker Analysis Table\n\n"
        "| Biomarker          | Your Value | Status | What This Means              | Normal Range                                 |\n"
        "|--------------------|------------|--------|-----------------------------|----------------------------------------------|\n"
        "Include ALL biomarkers: Albumin, Creatinine, Glucose, CRP, MCV, RDW, ALP, WBC, Lymphocytes, Hemoglobin, Plasma Volume\n"
        "Use emojis in Status column: 🟢 Normal, 🟡 Borderline, 🔴 High/Low\n\n"
        
        "---\n\n"
        "## 📋 Quick Summary\n\n"
        "### 🎯 Top 3 Health Priorities\n"
        "1. **[Priority 1]** - Brief explanation\n"
        "2. **[Priority 2]** - Brief explanation\n"
        "3. **[Priority 3]** - Brief explanation\n\n"
        
        "### ✅ What's Working Well\n"
        "• List normal biomarkers and what they indicate\n\n"
        
        "---\n\n"
        "## 🏥 Detailed System Analysis\n\n"
        "For each system, use this format:\n"
        "### [Emoji] [System Name] - Status: [🟢/🟡/🔴]\n"
        "Brief explanation in simple language.\n\n"
        "Cover: Metabolic, Kidney, Blood Health, Immune System, Circulation\n\n"
        
        "---\n\n"
        "## 💡 Your Action Plan\n\n"
        "### 🥗 Nutrition Recommendations\n"
        "• Specific dietary advice\n\n"
        "### 🏃 Lifestyle Changes\n"
        "• Exercise and daily habits\n\n"
        "### 🔍 Recommended Tests\n"
        "• Follow-up tests to discuss with your doctor\n\n"
        "### 👨‍⚕️ Medical Consultation\n"
        "• When and why to see your healthcare provider\n\n"
        
        "---\n\n"
        "## ⚠️ Important Connections\n\n"
        "Highlight how different biomarkers relate to each other.\n\n"
        
        "---\n\n"
        "## 📌 Important Reminders\n\n"
        "• This is an AI analysis, not a medical diagnosis\n"
        "• Always consult your healthcare provider\n"
        "• Bring this report to your next appointment\n\n"
        
        "NORMAL RANGE DEFINITIONS:\n"
        "• **Glucose (fasting)**: 70–99 mg/dL → <70: Low, 70–99: Normal, >99: High\n"
        "• **Albumin**: 3.4–5.4 g/dL → <3.4: Low, 3.4–5.4: Normal, >5.4: High\n"
        "• **Creatinine**: 0.6–1.3 mg/dL → <0.6: Low, 0.6–1.3: Normal, >1.3: High\n"
        "• **RDW**: 11.5–14.5% → <11.5: Low, 11.5–14.5: Normal, >14.5: High\n"
        "• **Alkaline Phosphatase (ALP)**: 44–147 IU/L → <44: Low, 44–147: Normal, >147: High\n"
        "• **WBC Count**: 4.0–11.0 ×10⁹/L → <4.0: Low, 4.0–11.0: Normal, >11.0: High\n"
        "• **CRP**: <5 mg/L → <1: Low, 1–5: Normal, >5: High\n"
        "• **MCV**: 80–100 fL → <80: Low, 80–100: Normal, >100: High\n"
        "• **Hemoglobin**: 13.8–17.2 g/dL (men), 12.1–15.1 g/dL (women) → Below or above these ranges: Low or High\n"
        "• **Plasma Volume**: Approximately 50–55 mL/kg body weight → Values outside this range may indicate fluid imbalances\n"
        "• **Lymphocytes**: 1.0–3.0 ×10⁹/L → <1.0: Low, 1.0–3.0: Normal, >3.0: High\n"
        "• Units must be shown clearly in reports\n\n"
        
        "CONSTRAINTS:\n"
        "• Never diagnose or prescribe medication\n"
        "• Use simple, encouraging language\n"
        "• Include reference ranges for all biomarkers\n"
        "• Make the report visually scannable with emojis and formatting"
    )

    # ----- User Message -----
    user_message = (
        f"Patient Information:\n"
        f"- Age: {age} years\n"
        f"- Gender: {gender}\n"
        f"- Height: {height} cm\n"
        f"- Weight: {weight} kg\n\n"
        f"Biomarker Values:\n"
        f"- Albumin: {albumin} g/dL\n"
        f"- Creatinine: {creatinine} mg/dL\n"
        f"- Glucose: {glucose} mg/dL\n"
        f"- CRP: {crp} mg/L\n"
        f"- MCV: {mcv} fL\n"
        f"- RDW: {rdw} %\n"
        f"- ALP: {alp} U/L\n"
        f"- WBC: {wbc} x10^3/μL\n"
        f"- Lymphocytes: {lymphocytes} %\n"
        f"- Hemoglobin: {hemoglobin} g/dL\n"
        f"- Plasma(PV) (ML): {pv} ML"
    )

    # ----- Generate AI Response -----
    try:
        model = genai.GenerativeModel(MODEL_ID)
        prompt = f"{system_message}\n\n{user_message}"
        
        response = model.generate_content(prompt)
        return response.text
    except Exception as e:
        return f"Error: {str(e)}\n\nPlease check the model name. Try running the list_models.py script first."


# ---------------- Gradio UI ----------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🏥 AI Health Assistant
        ### Powered by Gemini 2.0
        Get instant insights from your biomarker results with AI-powered analysis.
        """
    )

    with gr.Row():
        with gr.Column():
            albumin = gr.Textbox(label="Albumin (g/dL)", value="4.5")
            creatinine = gr.Textbox(label="Creatinine (mg/dL)", value="1.5")
            glucose = gr.Textbox(label="Glucose (mg/dL, fasting)", value="160")
            crp = gr.Textbox(label="CRP (mg/L)", value="2.5")
            mcv = gr.Textbox(label="MCV (fL)", value="150")
            rdw = gr.Textbox(label="RDW (%)", value="15")
            alp = gr.Textbox(label="ALP (U/L)", value="146")
            wbc = gr.Textbox(label="WBC (10^3/μL)", value="10.5")
            lymphocytes = gr.Textbox(label="Lymphocytes (%)", value="38")
            hemoglobin = gr.Textbox(label="Hemoglobin (g/dL)", value="13.5")
            pv = gr.Textbox(label="Plasma(PV) (ML)", value="3000")

        with gr.Column():
            age = gr.Textbox(label="Age (years)", value="30")
            gender = gr.Dropdown(choices=["Male", "Female"], label="Gender", value="Male")
            height = gr.Textbox(label="Height (cm)", value="170")
            weight = gr.Textbox(label="Weight (kg)", value="65")

    output = gr.Markdown(label="AI Health Report")

    btn = gr.Button("Generate Report")
    btn.click(
        respond,
        inputs=[
            albumin, creatinine, glucose, crp, mcv, rdw, alp, wbc,
            lymphocytes, hemoglobin, pv, age, gender, height, weight
        ],
        outputs=output
    )

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