# app.py — HealthSense AI
# Deployed on Hugging Face Spaces (Gradio SDK)

import json
import os
import re

import gradio as gr
import joblib
import spaces  # ADDED: required for ZeroGPU runtime contract (was live on Space, missing here)
from groq import Groq
from rapidfuzz import fuzz, process

# --- Load trained ML artifacts ---
MODEL_DIR = "models"
model = joblib.load(os.path.join(MODEL_DIR, "disease_model.joblib"))
encoder = joblib.load(os.path.join(MODEL_DIR, "label_encoder.joblib"))

with open(os.path.join(MODEL_DIR, "symptoms.json")) as f:
    symptoms_list = json.load(f)

# --- Groq client (reads GROQ_API_KEY from Space secrets) ---
groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# --- Voice-to-symptom matching setup ---
# Readable ("high fever") <-> dataset key ("high_fever") lookup tables
SYMPTOM_READABLE = {s: s.replace("_", " ").strip() for s in symptoms_list}
READABLE_TO_SYMPTOM = {v: k for k, v in SYMPTOM_READABLE.items()}
READABLE_CHOICES = list(SYMPTOM_READABLE.values())

# Everyday phrasing that has no character overlap with the clinical term,
# so fuzzy string matching alone would miss it. Extend this list over time
# based on what real users say.
SYMPTOM_SYNONYMS = {
    "throwing up": "vomiting",
    "throw up": "vomiting",
    "can't stop sneezing": "continuous_sneezing",
    "stuffy nose": "congestion",
    "tired": "fatigue",
    "exhausted": "fatigue",
    "dizzy": "dizziness",
    "upset stomach": "stomach_pain",
    "sore throat": "throat_irritation",
    "throat hurts": "throat_irritation",
}


def transcribe_audio(audio_filepath):
    """Send recorded audio to Groq's hosted Whisper model and return the text."""
    with open(audio_filepath, "rb") as f:
        transcription = groq_client.audio.transcriptions.create(
            file=f,
            model="whisper-large-v3-turbo",
        )
    return transcription.text


def match_symptoms_from_text(transcript, score_cutoff=68):
    """Map free-form English text onto the fixed 132-item symptom checklist.
    Fast, local, zero API cost — but only works for Latin-script/English text.
    Used as a fallback if the LLM extractor below fails."""
    text = transcript.lower()
    matched = set()

    # Pass 1: known synonym phrases
    for phrase, canonical in SYMPTOM_SYNONYMS.items():
        if phrase in text:
            matched.add(canonical)

    # Pass 2: fuzzy match each clause of the transcript against the vocabulary
    clauses = re.split(r",| and | since | for | with ", text)
    for clause in clauses:
        clause = clause.strip()
        if len(clause) < 3:
            continue
        result = process.extractOne(
            clause, READABLE_CHOICES, scorer=fuzz.token_set_ratio, score_cutoff=score_cutoff
        )
        if result:
            matched_readable = result[0]
            matched.add(READABLE_TO_SYMPTOM[matched_readable])

    return sorted(matched)


def extract_symptoms_llm(user_text):
    """Map free-form text in ANY language (Hindi, English, mixed) onto the
    fixed symptom list using the LLM itself — this is what makes multilingual
    input work, since string-matching cannot bridge Devanagari and Latin script."""
    symptom_options = ", ".join(symptoms_list)
    prompt = f"""You are a medical symptom extraction assistant. A patient described how \
they feel, possibly in Hindi, English, or a mix of both.

Patient description: "{user_text}"

From this exact list of allowed symptoms, return ONLY the ones the patient is \
describing, as a JSON array of exact strings copied from the list below. Do not \
invent new symptom names, do not translate or reword them — only use exact matches \
from this list:
{symptom_options}

Respond with ONLY a JSON array, nothing else. Example: ["headache", "high_fever", "nausea"]"""

    try:
        response = groq_client.chat.completions.create(
            model="openai/gpt-oss-20b",
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
        )
        raw = response.choices[0].message.content.strip()
        if raw.startswith("```"):
            raw = raw.strip("`").replace("json", "", 1).strip()
        matched = json.loads(raw)
        # Safety net: only keep symptoms that actually exist in our vocabulary
        matched = sorted(set(s for s in matched if s in symptoms_list))
        return matched
    except Exception:
        # LLM call failed or returned unparseable output — fall back to local
        # fuzzy matching (English-only, but better than nothing)
        return match_symptoms_from_text(user_text)


def predict_disease(selected_symptoms):
    """Run the Random Forest classifier on the selected symptoms."""
    input_vec = [[1 if s in selected_symptoms else 0 for s in symptoms_list]]
    proba = model.predict_proba(input_vec)[0]
    top_idx = int(proba.argmax())
    disease = encoder.inverse_transform([top_idx])[0]
    confidence = round(float(proba[top_idx]) * 100, 1)
    return disease, confidence


def get_ai_guidance(disease, symptoms, confidence):
    """Ask the LLM for a bilingual (Hindi-primary), voice-ready explanation,
    with a trailing [AUDIO_PLAY_TEXT: ...] tag for the 🔊 speaker button."""
    prompt = f"""You are a bilingual Medical & Health Information Assistant.

Patient symptoms: {', '.join(symptoms)}.
ML model predicted: {disease} ({confidence}% confidence).

Respond following these rules exactly:

1. DUAL-LANGUAGE FORMAT: Write the entire explanation in simple, easy-to-understand \
Hindi (Devanagari script). Whenever you use an English medical term or medicine name, \
explain it in simple Hindi immediately after, in brackets, so a non-technical reader \
understands it without help.

2. VOICE & AUDIO READY: Keep the tone conversational and empathetic — natural to listen \
to if read aloud. Use clear bullet points for scannability.

3. Structure with these sections, each with a short Hindi heading:
   - रोग के बारे में (About the condition)
   - लक्षणों की व्याख्या (Your symptoms explained)
   - घरेलू देखभाल (Home care — 3 bullet points)
   - सामान्य दवाइयाँ (General over-the-counter medicine categories only — no specific \
     dosages; always say to confirm the dose with a pharmacist or doctor)
   - कब डॉक्टर से मिलें (When to seek urgent care)
   - बचाव के उपाय (Prevention tips)

4. SAFETY: End the main explanation with a brief note recommending consultation with a \
qualified doctor.

5. Immediately after that, on its own final line, add exactly one tag in this format \
(nothing after it): [AUDIO_PLAY_TEXT: <a 2-3 sentence Hindi summary of the condition and \
the single most important next step, written for text-to-speech>]"""

    try:
        response = groq_client.chat.completions.create(
            model="openai/gpt-oss-20b",
            messages=[{"role": "user", "content": prompt}],
        )
        return response.choices[0].message.content
    except Exception as e:
        return (
            "⚠️ Couldn't reach the AI guidance service right now "
            f"({e}). The ML prediction above is still valid."
        )


def parse_audio_tag(guidance_text):
    """Split the LLM output into (clean_display_text, tts_summary_text).
    The [AUDIO_PLAY_TEXT: ...] tag is stripped from what's shown on screen
    and instead bound to the 🔊 speaker button for text-to-speech."""
    match = re.search(r"\[AUDIO_PLAY_TEXT:\s*(.*?)\]\s*$", guidance_text, re.DOTALL)
    if match:
        clean_text = guidance_text[: match.start()].strip()
        audio_text = match.group(1).strip()
        return clean_text, audio_text
    return guidance_text, ""


@spaces.GPU  # ADDED: ZeroGPU runtime requires this on the function that's actually
             # invoked when symptoms are submitted, not a dummy/unused function.
             # App itself is CPU-only; this satisfies HF's ZeroGPU startup contract.
def analyse(selected_symptoms):
    if not selected_symptoms:
        return "⚠️ Please select at least one symptom.", "", ""

    disease, confidence = predict_disease(selected_symptoms)
    raw_guidance = get_ai_guidance(disease, selected_symptoms, confidence)
    guidance, audio_text = parse_audio_tag(raw_guidance)  # ADDED: split off TTS tag
    result = f"### 🩺 Predicted: **{disease}** — {confidence}% confidence"
    return result, guidance, audio_text


def pipeline_from_text(user_text, source_label="🗨️ You said"):
    """Shared core: raw text (typed OR transcribed) -> matched symptoms ->
    prediction -> guidance. Both the text box and the voice input funnel
    through this same function."""
    if not user_text or not user_text.strip():
        return [], "", "", "", "", ""

    matched_symptoms = extract_symptoms_llm(user_text)

    if not matched_symptoms:
        transcript_display = (
            f"{source_label}: \"{user_text}\"\n\n"
            "⚠️ Couldn't confidently match this to any known symptoms — "
            "please try rephrasing or select manually below."
        )
        return [], transcript_display, "", "", "", ""

    ml_result, guidance, audio_text = analyse(matched_symptoms)  # FIX: 3 return values now
    transcript_display = (
        f"{source_label}: \"{user_text}\"\n\n"
        f"✅ Matched symptoms: {', '.join(s.replace('_', ' ') for s in matched_symptoms)}"
    )
    # Last output clears the textbox after submission, for a chat-like feel
    return matched_symptoms, transcript_display, ml_result, guidance, audio_text, ""


def text_pipeline(user_text):
    return pipeline_from_text(user_text, source_label="🗨️ You typed")


def voice_pipeline(audio_filepath):
    """Full automated flow: transcribe -> extract symptoms -> predict -> guidance."""
    if audio_filepath is None:
        return [], "🎙️ No audio received — please try recording again.", "", "", "", ""

    transcript = transcribe_audio(audio_filepath)
    matched, display, ml_result, guidance, audio_text, _ = pipeline_from_text(
        transcript, source_label="🎙️ You said"
    )
    return matched, display, ml_result, guidance, audio_text, ""


DISCLAIMER = (
    "⚠️ **Disclaimer:** HealthSense AI is an educational demo, not a substitute "
    "for professional medical advice. Always consult a qualified doctor for "
    "diagnosis and treatment."
)

with gr.Blocks(title="HealthSense AI", theme=gr.themes.Soft()) as app:
    gr.Markdown(
        "# 🏥 HealthSense AI\n"
        "Type how you feel or speak into the mic — in Hindi or English — "
        "and get an instant AI-powered prediction."
    )

    # --- Top-of-page medical disclaimer badge (visible before any interaction) ---
    gr.Markdown(
        "> ⚠️ **Educational tool only — not a medical diagnosis.** / "
        "**केवल एक शैक्षणिक उपकरण — चिकित्सा निदान नहीं।** "
        "Always consult a qualified doctor. / हमेशा किसी योग्य डॉक्टर से सलाह लें।"
    )

    # --- Onboarding: dual-language getting-started guide ---
    with gr.Accordion("ℹ️ How to use / उपयोग कैसे करें", open=True):
        gr.Markdown(
            "**🇬🇧 English**\n"
            "1. Type your symptoms below, **or** tap the 🎙️ mic and speak naturally.\n"
            "2. First time using voice? Your browser will ask for microphone "
            "permission — tap **Allow**. Nothing is recorded without this.\n"
            "3. Speak or type clearly, e.g. *\"I have a headache and fever since "
            "yesterday.\"*\n"
            "4. Symptoms are matched and analysed automatically — no extra clicks.\n"
            "5. Read your result below, and always confirm with a real doctor.\n\n"
            "**🇮🇳 हिंदी**\n"
            "1. नीचे अपने लक्षण टाइप करें, **या** 🎙️ माइक्रोफ़ोन पर टैप करके "
            "स्वाभाविक रूप से बोलें।\n"
            "2. पहली बार आवाज़ का उपयोग कर रहे हैं? ब्राउज़र माइक्रोफ़ोन अनुमति "
            "मांगेगा — **Allow** पर टैप करें। इसके बिना कुछ भी रिकॉर्ड नहीं होता।\n"
            "3. स्पष्ट रूप से बोलें या टाइप करें, जैसे *\"मुझे कल से सिरदर्द और "
            "बुखार है।\"*\n"
            "4. लक्षणों का मिलान और विश्लेषण अपने आप हो जाता है — किसी अतिरिक्त "
            "क्लिक की ज़रूरत नहीं।\n"
            "5. नीचे अपना परिणाम पढ़ें, और हमेशा किसी वास्तविक डॉक्टर से पुष्टि करें।"
        )

    gr.Markdown("### 💬 Describe your symptoms")
    with gr.Row():
        text_input = gr.Textbox(
            placeholder="e.g. \"mujhe sar dard aur bukhar hai\" or \"I have a headache and fever\"",
            show_label=False,
            scale=4,
        )
        voice_input = gr.Audio(
            sources=["microphone"], type="filepath", label="🎙️ Tap to speak", scale=1
        )

    # --- Sample prompts so first-time users know what "good input" looks like ---
    gr.Examples(
        examples=[
            "I have a headache, high fever, and nausea since yesterday",
            "mujhe khansi aur gale mein dard hai",
            "My skin is itchy with red rashes and swelling",
            "पेट में दर्द और बार-बार उल्टी हो रही है",
        ],
        inputs=text_input,
        label="💡 Try an example / उदाहरण आज़माएं",
    )

    transcript_output = gr.Markdown(label="Recognition Result")

    with gr.Accordion("Or select symptoms manually", open=False):
        symptom_input = gr.CheckboxGroup(choices=symptoms_list, label="Select Symptoms")
        btn = gr.Button("🔍 Analyse Symptoms", variant="primary")

    ml_output = gr.Markdown(label="ML Prediction")
    ai_output = gr.Markdown(label="AI Medical Guidance")
    audio_text_state = gr.Textbox(visible=False)  # ADDED: holds TTS summary for speaker button
    speaker_btn = gr.Button("🔊 सुनें (Listen)")  # ADDED: reads the Hindi summary aloud

    gr.Markdown(DISCLAIMER)

    # Manual path: check boxes yourself, click the button
    btn.click(
        fn=analyse,
        inputs=symptom_input,
        outputs=[ml_output, ai_output, audio_text_state],  # FIX: 3 outputs now
    )

    # Text path: type + press Enter -> fully automated
    text_input.submit(
        fn=text_pipeline,
        inputs=text_input,
        outputs=[symptom_input, transcript_output, ml_output, ai_output, audio_text_state, text_input],
    )

    # Voice path: recording stops -> fully automated
    voice_input.stop_recording(
        fn=voice_pipeline,
        inputs=voice_input,
        outputs=[symptom_input, transcript_output, ml_output, ai_output, audio_text_state, text_input],
    )

    # ADDED: Speaker button -> browser-native speechSynthesis (free, zero API cost,
    # matches the lightweight TTS recommendation from the project's voice roadmap).
    # No server round-trip: the browser itself reads audio_text_state's current value.
    speaker_btn.click(
        fn=None,
        inputs=[audio_text_state],
        outputs=[],
        js="""(text) => {
            if (!text) return;
            window.speechSynthesis.cancel();
            const utterance = new SpeechSynthesisUtterance(text);
            utterance.lang = 'hi-IN';
            window.speechSynthesis.speak(utterance);
        }""",
    )

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