"""
Medical AI Assistant – Pollinations GET flavour
Author: you
"""
import gradio as gr
import requests
import urllib.parse
import json
import base64
from io import BytesIO
from docx import Document

# ------------------------------------------------------------------
# 1.  Generic Pollinations helper (GET /encoded_prompt)
# ------------------------------------------------------------------
def get_pollinations_response(
        prompt: str,
        model: str = "openai",
        seed: int = 42,
        system_prompt: str = ""
) -> str:
    """
    Calls https://text.pollinations.ai/{encoded_prompt}
    Returns plain-text string (or error string).
    """
    params = {"model": model, "seed": seed}
    if system_prompt:
        params["system"] = system_prompt

    encoded = urllib.parse.quote(prompt)
    url = f"https://text.pollinations.ai/{encoded}"

    try:
        resp = requests.get(url, params=params, timeout=60)
        resp.raise_for_status()
        return resp.text
    except Exception as e:
        return f"API error: {e}"

# ------------------------------------------------------------------
# 2.  Medical-specific wrappers
# ------------------------------------------------------------------
def generate_diagnosis(symptoms, history, age, gender, allergies, meds, family, lifestyle):
    system = ("You are an expert medical diagnostician. "
              "Provide evidence-based PRELIMINARY diagnoses ranked by probability.")
    user = f"""
Patient: {age} y/o {gender}
Symptoms: {symptoms}
History: {history}
Allergies: {allergies or 'None'}
Meds: {meds or 'None'}
Family: {family or 'None'}
Lifestyle: {lifestyle or 'None'}

Give:
1. Most likely conditions (ranked)
2. Severity
3. Clinical reasoning"""
    return get_pollinations_response(user, system_prompt=system)

def generate_treatment_plan(symptoms, history, age, gender, allergies, meds, family, diagnosis):
    system = ("You are a specialist in personalised treatment. "
              "Suggest safe pharmacological + non-pharmacological steps.")
    user = f"""
Patient: {age} y/o {gender}
Diagnosis: {diagnosis}
History: {history}
Allergies: {allergies or 'None'}
Meds: {meds or 'None'}

Provide:
1. Medication & dosing
2. Lifestyle / diet
3. Follow-up timing
4. Red-flag symptoms
5. Rationale"""
    return get_pollinations_response(user, system_prompt=system)

# ------------------------------------------------------------------
# 3.  Word report builder
# ------------------------------------------------------------------
def build_docx(diagnosis, treatment, data):
    doc = Document()
    doc.add_heading("Healthcare AI Assistant Report", 0)
    doc.add_heading("Patient info", 1)
    doc.add_paragraph(f"Age: {data['age']}")
    doc.add_paragraph(f"Gender: {data['gender']}")
    doc.add_heading("Preliminary diagnosis", 1)
    doc.add_paragraph(diagnosis)
    doc.add_heading("Treatment plan", 1)
    doc.add_paragraph(treatment)
    doc.add_heading("Disclaimer", 1)
    doc.add_paragraph("This is an AI-assisted preliminary analysis – NOT a substitute for professional medical consultation.")

    bio = BytesIO()
    doc.save(bio)
    bio.seek(0)
    return bio

# ------------------------------------------------------------------
# 4.  Main orchestrator
# ------------------------------------------------------------------
def process_medical_analysis(symptoms, history, age, gender, allergies, meds, family, lifestyle):
    if not symptoms or not history:
        return "⚠️  Please provide both symptoms and medical history.", "", ""
    diagnosis = generate_diagnosis(symptoms, history, age, gender, allergies, meds, family, lifestyle)
    treatment = generate_treatment_plan(symptoms, history, age, gender, allergies, meds, family, diagnosis)
    docx_bio = build_docx(diagnosis, treatment, {"age": age, "gender": gender})
    b64 = base64.b64encode(docx_bio.read()).decode()
    link = f'<a href="data:application/vnd.openxmlformats-officedocument.wordprocessingml.document;base64,{b64}" download="medical_report.docx">📥 Download Report</a>'
    return diagnosis, treatment, link

# ------------------------------------------------------------------
# 5.  Gradio UI
# ------------------------------------------------------------------
with gr.Blocks(title="Medical AI Assistant", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🏥 Medical AI Assistant")
    gr.Markdown("AI-powered preliminary diagnosis & treatment plan using Pollinations AI")

    with gr.Row():
        with gr.Column():
            age = gr.Slider(0, 120, value=30, step=1, label="Age")
            gender = gr.Radio(["Male", "Female", "Other"], value="Male", label="Gender")
        with gr.Column():
            symptoms = gr.Textbox(label="Current symptoms", placeholder="e.g. fever 3 days, dry cough", lines=4)
            history = gr.Textbox(label="Medical history", placeholder="e.g. hypertension, diabetes 2019", lines=4)

    with gr.Row():
        with gr.Column():
            allergies = gr.Textbox(label="Known allergies", placeholder="e.g. penicillin")
            meds = gr.Textbox(label="Current medications", placeholder="e.g. metformin 500 mg bid")
        with gr.Column():
            family = gr.Textbox(label="Family history", placeholder="e.g. heart disease")
            lifestyle = gr.Textbox(label="Lifestyle", placeholder="e.g. smoker, exercises 3×/wk")

    go = gr.Button("🔍 Generate Analysis", variant="primary", size="lg")

    with gr.Row():
        diag_out = gr.Textbox(label="Preliminary Diagnosis", lines=10, interactive=False)
        treat_out = gr.Textbox(label="Treatment Plan", lines=10, interactive=False)
    download_html = gr.HTML()

    go.click(process_medical_analysis,
             inputs=[symptoms, history, age, gender, allergies, meds, family, lifestyle],
             outputs=[diag_out, treat_out, download_html])

    gr.Markdown("⚠️ **Disclaimer**: This tool provides preliminary AI-generated suggestions only – always consult a licensed healthcare professional.")

if __name__ == "__main__":
    demo.launch(share=True, server_name="0.0.0.0", server_port=7860)