import gradio as gr
from groq import Groq
import os

# 🔹 Set your Groq API Key securely
os.environ["GROQ_API_KEY"] = "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"

# 🔹 Initialize Groq client
client = Groq(api_key=os.getenv("GROQ_API_KEY"))

# 🔹 Define model
MODEL_ID = "llama-3.3-70b-versatile"

# ---------------- AI Response Function ----------------
def respond(albumin, creatinine, glucose, crp, mcv, rdw, alp, wbc, lymphocytes,
            hemoglobin, pv, age, gender, height, weight):
    
    # ----- System Prompt -----
    system_message = (
        "You are an AI Health Assistant that analyzes laboratory biomarkers "
        "and generates structured, patient-friendly health summaries.\n\n"
        "Your task is to evaluate the provided biomarkers and generate an AI-driven medical report "
        "with insights, observations, and clear explanations.\n"
        "You must strictly follow this structured format:\n\n"
        "### Tabular Mapping\n"
        "- Always include a Markdown table with exactly four columns:\n"
        "| Biomarker | Value | Status (Low/Normal/High) | AI-Inferred Insight | Reference Range |\n"
        "- Include **all available biomarkers** below:\n"
        "Albumin, Creatinine, Glucose, CRP, MCV, RDW, ALP, WBC, Lymphocytes, Hemoglobin, Plasma Viscosity (PV)\n"
        "- The first row after the header must begin directly with 'Albumin'.\n"
        "- Each biomarker must appear exactly once as a separate row.\n\n"
        "### Executive Summary\n"
        "- List Top 3 Health Priorities.\n"
        "- Highlight Key Strengths or normal biomarkers.\n\n"
        "### System-Specific Analysis\n"
        "- Summarize findings grouped by organ systems (Liver, Kidney, Immune, Blood, etc.).\n"
        "- Status: “Optimal” | “Monitor” | “Needs Attention”.\n"
        "- Provide 2–3 sentences of explanation in plain, supportive language.\n\n"
        "### Personalized Action Plan\n"
        "- Provide categorized recommendations (Nutrition, Lifestyle, Testing, Medical Consultation).\n"
        "- Never recommend medication or treatment.\n\n"
        "### Interaction Alerts\n"
        "- Highlight potential relationships between markers (e.g., high CRP + low Albumin).\n\n"
        "### Constraints\n"
        "- Never give a diagnosis or prescribe medicine.\n"
        "- Never use data not present in the input.\n"
        "- Always recommend consulting a healthcare professional.\n"
        "- Always include normal reference ranges for each biomarker.\n"
        "- Use simple, clear, patient-friendly language."
    )

    # ----- 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 Viscosity (PV): {pv} mPa·s"
    )

    # ----- Call Groq API -----
    completion = client.chat.completions.create(
        model=MODEL_ID,
        messages=[
            {"role": "system", "content": system_message},
            {"role": "user", "content": user_message}
        ],
        temperature=0.4,
        max_tokens=2500,
        top_p=0.9,
        stream=False
    )

    return completion.choices[0].message.content


# ---------------- Gradio UI ----------------
with gr.Blocks() as demo:
    gr.Markdown("## 🧪 AI Health Assistant (Extended Biomarkers via Groq Llama-3.3-70B)")

    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 Viscosity (mPa·s)", value="1.7")

        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.Textbox(label="AI Health Report", lines=30)

    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()