"""
Medical Image Triage — HuggingFace Space (CPU)
Model : Qwen/Qwen2-VL-2B-Instruct  (transformers, CPU inference)
Memory: ChromaDB + all-MiniLM-L6-v2 embeddings
UI    : Gradio
"""

import hashlib
import logging
import os

import chromadb
import gradio as gr
import torch
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
from huggingface_hub import login
from PIL import Image
from qwen_vl_utils import process_vision_info
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration

# ── Logging ────────────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

# ── HuggingFace auth ───────────────────────────────────────────────────────────
HF_TOKEN = os.environ.get("HF_TOKEN", "")
if HF_TOKEN:
    login(token=HF_TOKEN)
    log.info("Logged in to HuggingFace Hub.")
else:
    log.warning("HF_TOKEN secret not set — model download may fail for gated repos.")

# ── Model config ───────────────────────────────────────────────────────────────
# 2B is the practical limit for CPU inference; 7B would take many minutes per image.
DEFAULT_MODEL_ID = "Qwen/Qwen2-VL-2B-Instruct"

# ── Vector DB ──────────────────────────────────────────────────────────────────
VECTOR_DB_PATH  = os.path.join(os.getcwd(), "medical_memory_chroma")
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

log.info("Initialising ChromaDB …")
_embed_fn = SentenceTransformerEmbeddingFunction(model_name=EMBEDDING_MODEL)
_chroma   = chromadb.PersistentClient(path=VECTOR_DB_PATH)
medical_collection = _chroma.get_or_create_collection(
    name="medical_triage_notes",
    embedding_function=_embed_fn,
)
log.info("ChromaDB ready.")


# ── Inference class ────────────────────────────────────────────────────────────
class ImageInference:
    """Qwen2-VL vision-language model running on CPU via transformers."""

    def __init__(self, model_name: str = DEFAULT_MODEL_ID):
        log.info("Loading model: %s  (CPU — this takes a minute) …", model_name)
        self.model_name = model_name

        self.model = Qwen2VLForConditionalGeneration.from_pretrained(
            model_name,
            torch_dtype=torch.float32,   # float32 for CPU stability
            device_map="cpu",
        )
        self.model.eval()

        self.processor = AutoProcessor.from_pretrained(
            model_name, trust_remote_code=True
        )
        log.info("Model ready: %s", model_name)

    def generate_image_output(
        self, image: Image.Image, patient_context: str = ""
    ) -> str:
        context_block = (
            f"Patient context: {patient_context}\n" if patient_context.strip() else ""
        )

        triage_prompt = (
            "You are a medical image triage assistant. "
            "Analyze the provided image and return a concise structured assessment.\n"
            "Classify the image as one of: xray, normal_photo, prescription, or unknown.\n"
            "If the image looks like a prescription, extract the visible text exactly.\n"
            "If the image looks like a medical photo or X-ray, give a conservative "
            "triage label: normal, monitor, urgent, or emergency.\n"
            "Use the following format exactly:\n"
            "image_type: <xray|normal_photo|prescription|unknown>\n"
            "triage_label: <normal|monitor|urgent|emergency|not_applicable>\n"
            "summary: <one short sentence>\n"
            "findings: <bullet-style semicolon-separated details>\n"
            "prescription_text: <exact text or none>\n"
            "follow_up_questions: <up to 3 questions, comma-separated>\n"
            f"{context_block}"
            "Do not provide a final diagnosis. "
            "Do not add commentary outside the requested format."
        )

        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text",  "text": triage_prompt},
                ],
            }
        ]

        text_input = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        image_inputs, video_inputs = process_vision_info(messages)

        inputs = self.processor(
            text=[text_input],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )

        with torch.no_grad():
            generated_ids = self.model.generate(
                **inputs,
                max_new_tokens=512,
                do_sample=False,          # greedy — faster on CPU
                temperature=None,
                top_p=None,
            )

        # Strip the prompt tokens from the output
        generated_ids_trimmed = [
            out_ids[len(in_ids):]
            for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]
        output_text = self.processor.batch_decode(
            generated_ids_trimmed,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )
        return output_text[0] if output_text else "No output generated."


# ── Load model at startup ──────────────────────────────────────────────────────
log.info("Loading model at startup …")
inference     = ImageInference(DEFAULT_MODEL_ID)
current_model = DEFAULT_MODEL_ID


# ── Utilities ──────────────────────────────────────────────────────────────────
def triage_text_to_dict(text: str) -> dict:
    out = {}
    for line in text.splitlines():
        line = line.strip()
        if not line or ":" not in line:
            continue
        k, v = line.split(":", 1)
        out[k.strip()] = v.strip()
    if "follow_up_questions" in out:
        out["follow_up_questions"] = [
            q.strip()
            for q in out["follow_up_questions"].split(",")
            if q.strip()
        ]
    return out


def upsert_triage_to_chroma(triage_report: dict, conversation_id: str = "default") -> str:
    document = "\n".join([
        f"image_type: {triage_report.get('image_type', '')}",
        f"triage_label: {triage_report.get('triage_label', '')}",
        f"summary: {triage_report.get('summary', '')}",
        f"findings: {triage_report.get('findings', '')}",
        f"prescription_text: {triage_report.get('prescription_text', 'none')}",
        f"follow_up_questions: {', '.join(triage_report.get('follow_up_questions', []))}",
    ])
    record_id = hashlib.sha1(f"{conversation_id}:{document}".encode()).hexdigest()
    medical_collection.upsert(
        ids=[record_id],
        documents=[document],
        metadatas=[{"conversation_id": conversation_id, "kind": "triage_report"}],
    )
    return record_id


# ── Gradio callbacks ───────────────────────────────────────────────────────────
def analyze_image(image: Image.Image, patient_context: str):
    if image is None:
        return "⚠️ Please upload an image first."
    try:
        pil_image   = image.convert("RGB")
        result_text = inference.generate_image_output(
            pil_image, patient_context=patient_context or ""
        )
        triage_report = triage_text_to_dict(result_text)

        record_id = None
        try:
            record_id = upsert_triage_to_chroma(triage_report)
        except Exception as exc:
            log.warning("ChromaDB upsert failed: %s", exc)

        label    = triage_report.get("triage_label", "—").upper()
        img_type = triage_report.get("image_type",   "—")
        summary  = triage_report.get("summary",      "—")
        findings = triage_report.get("findings",     "—")
        rx_text  = triage_report.get("prescription_text", "none")
        follow_ups = triage_report.get("follow_up_questions", [])

        badge = {"NORMAL": "🟢", "MONITOR": "🟡", "URGENT": "🟠", "EMERGENCY": "🔴"}.get(label, "⚪")
        follow_up_md = "\n".join(f"- {q}" for q in follow_ups) if follow_ups else "—"

        return f"""
## {badge} Triage Report

| Field | Value |
|---|---|
| **Image Type** | {img_type} |
| **Triage Label** | {label} |
| **Summary** | {summary} |

### 🔍 Findings
{findings}

### 💊 Prescription Text
{rx_text}

### ❓ Follow-up Questions
{follow_up_md}

---
*Record stored in vector DB: `{record_id or 'N/A'}`*
""".strip()

    except Exception as exc:
        log.exception("analyze_image error")
        return f"❌ Error: {exc}"


# ── Gradio UI ──────────────────────────────────────────────────────────────────
CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@400;600;700&display=swap');
.title-box {
    text-align: center; border: 2px solid #d1d5db; border-radius: 14px;
    padding: 20px; margin-bottom: 20px;
    background: linear-gradient(135deg,#f0f9ff 0%,#e0f2fe 100%);
    font-family: 'Space Grotesk', sans-serif;
}
.title-box h1 { margin-bottom: 8px; font-size: 38px; font-weight: 700; }
.title-box p  { font-size: 15px; color: #4b5563; }
"""

with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:

    gr.Markdown("""
<div class="title-box">
    <h1>🏥 Dr. ROCM</h1>
    <p>Upload an X-ray, clinical photo, or prescription.<br>
    The model returns a structured triage report.<br>
    <small>⚠️ Running on CPU — inference takes ~60 seconds per image.</small></p>
</div>
""")

    with gr.Row():
        with gr.Column(scale=1):
            image_input   = gr.Image(type="pil", label="Upload Image")
            context_input = gr.Textbox(
                label="Patient Context (optional)",
                placeholder="e.g. 45-year-old male, chest pain for 2 days …",
                lines=3,
            )
            analyze_btn = gr.Button("🔍 Run Triage Analysis", variant="primary")

        with gr.Column(scale=1):
            output_markdown = gr.Markdown("### Results will appear here …")

    analyze_btn.click(
        analyze_image,
        inputs=[image_input, context_input],
        outputs=output_markdown,
    )

    gr.Markdown(
        "_⚠️ This tool is for **triage assistance only** and does not constitute "
        "a medical diagnosis. Always consult a qualified healthcare professional._"
    )

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