from utils import chatbot_prompt, report_prompt
from lab_logic import analyze_lab_json
from rag_engine import match_disease

import gradio as gr
import os
import json
from huggingface_hub import InferenceClient


# 🔐 Load HF Token
token = os.getenv("HF_TOKEN")

if token is None:
    raise ValueError("HF_TOKEN not found in environment variables")


# 🔥 Create client once
client = InferenceClient(
    model="Intelligent-Internet/II-Medical-8B",
    provider="featherless-ai",
    token=token,
    timeout=120
)


# =========================
# 🔥 LLM CALL
# =========================
def call_model(prompt, system_instruction, max_tokens=300):

    try:

        response = client.chat_completion(
            messages=[
                {"role": "system", "content": system_instruction},
                {"role": "user", "content": prompt}
            ],
            max_tokens=max_tokens,
            temperature=0.02
        )

        text = ""

        if response and response.choices:
            msg = response.choices[0].message

            # normal response
            if hasattr(msg, "content") and msg.content:
                text = msg.content

            # fallback (featherless models)
            elif hasattr(msg, "reasoning") and msg.reasoning:
                text = msg.reasoning

        if not text:
            print("MODEL RESPONSE EMPTY:", response)
            return ""

        text = text.replace("<think>", "").replace("</think>", "").strip()

        return text

    except Exception as e:
        print("MODEL ERROR:", e)
        return clean_response(text)

def clean_response(text):

    if not text:
        return ""

    # remove reasoning style sentences
    lines = text.split("\n")

    filtered = []

    for line in lines:

        line_lower = line.lower()

        if line_lower.startswith("okay"):
            continue

        if "the user" in line_lower:
            continue

        if "let's" in line_lower:
            continue

        if "i need to" in line_lower:
            continue

        filtered.append(line)

    return " ".join(filtered).strip()


# =========================
# 🩺 SYMPTOM CHECKER
# =========================
def run_symptom(symptoms):

    if not symptoms.strip():
        return "Please enter symptoms."

    symptoms_lower = symptoms.lower()

    emergency_keywords = [
        "chest pain",
        "severe chest pain",
        "breathing difficulty",
        "shortness of breath",
        "unable to breathe",
        "unconscious",
        "fainting",
        "loss of consciousness",
        "severe bleeding",
        "blood vomiting",
        "vomiting blood",
        "stroke",
        "paralysis",
        "heart attack",
        "seizure",
        "fits"
    ]

    if any(k in symptoms_lower for k in emergency_keywords):
        return """⚠️ Possible Medical Emergency Detected
Immediate medical attention is required.
Please go to the nearest hospital.

⚠️ संभावित आपातकालीन स्थिति
तुरंत चिकित्सा सहायता लें।
निकटतम अस्पताल जाएँ।
"""

    extraction_prompt = f"""
Extract key medical symptoms from the sentence.
Return only comma separated symptoms.

Input:
{symptoms}
"""

    key_symptoms = call_model(
        extraction_prompt,
        system_instruction="Return only symptoms separated by comma. No explanation.",
        max_tokens=40
    )

    if not key_symptoms:
        key_symptoms = symptoms

    matches = match_disease(key_symptoms)

    if not matches:
        return "No matching condition found."

    unique = {}

    for m in matches:
        disease = m["disease"]
        score = m["score"]

        if disease not in unique or score > unique[disease]:
            unique[disease] = score

    sorted_matches = sorted(unique.items(), key=lambda x: x[1], reverse=True)

    top3 = sorted_matches[:3]

    output = "Top Possible Conditions:\n\n"

    for i, (disease, score) in enumerate(top3, start=1):
        output += f"{i}. {disease} — {score}%\n"

    main_disease = top3[0][0]

    output += "\nEnglish Summary:\n"
    output += (
        f"Your symptoms ({key_symptoms}) most closely match {main_disease}. "
        "Other possible conditions are listed above. "
        "If symptoms continue or worsen, consult a healthcare professional.\n"
    )

    output += "\nHindi Summary:\n"
    output += (
        f"आपके लक्षण ({key_symptoms}) सबसे अधिक {main_disease} से मेल खाते हैं। "
        "ऊपर अन्य संभावित स्थितियाँ भी दिखाई गई हैं। "
        "यदि लक्षण बने रहते हैं या बढ़ते हैं तो डॉक्टर से सलाह लें।"
    )

    return output


# =========================
# 📄 REPORT ANALYZER
# =========================
def run_report(report_json_text):

    if not report_json_text.strip():
        return "No report data received."

    try:
        report_json = json.loads(report_json_text)
    except:
        return "Invalid JSON format."

    if "report_info" not in report_json:
        return "Invalid report structure."

    abnormal = analyze_lab_json(report_json)

    if not abnormal:
        return "All parameters are within normal range."

    patient_name = report_json.get("patient_info", {}).get("name", "Patient")

    prompt = report_prompt(patient_name, abnormal)

    result = call_model(
        prompt,
        system_instruction="Follow format exactly. No extra text.",
        max_tokens=300
    )

    if not result:
        return "Unable to analyze report right now."

    return result


# =========================
# 💬 CHATBOT
# =========================
def run_chat(message):

    if not message.strip():
        return "Please enter your question."

    prompt = chatbot_prompt(message)

    result = call_model(
        prompt,
        system_instruction="Respond like a medical doctor giving short advice.",
        max_tokens=120
    )

    if not result:
        return "AI could not generate a response."

    return result


# =========================
# 🎨 UI
# =========================
with gr.Blocks(fill_height=True) as demo:

    gr.Markdown("# 🩺 Sehat Smartcare AI")

    with gr.Tabs():

        with gr.Tab("Symptom Checker"):
            symptom_input = gr.Textbox(label="Enter your symptoms")
            symptom_btn = gr.Button("Analyze")
            symptom_output = gr.Textbox(label="Result", lines=18)
            symptom_btn.click(run_symptom, symptom_input, symptom_output)

        with gr.Tab("Report Analyzer (JSON Input)"):
            report_input = gr.Textbox(label="Paste Structured JSON Report", lines=15)
            report_btn = gr.Button("Analyze Report")
            report_output = gr.Textbox(label="Clinical Summary", lines=12)
            report_btn.click(run_report, report_input, report_output)

        with gr.Tab("Medical Chatbot"):
            chat_input = gr.Textbox(label="Ask your health question", lines=4)
            chat_btn = gr.Button("Send")
            chat_output = gr.Textbox(label="Response", lines=8)
            chat_btn.click(run_chat, chat_input, chat_output)

demo.launch()