"""
llm_extractor.py — HuggingFace Spaces version.
Uses Phi-3 Mini (3.8B) running locally in the same Python process.
Patient data stays on HF Spaces server — does NOT leave to any external API.

Why Phi-3 Mini:
- 3.8B parameters, fits in HF free 16GB RAM
- Specifically strong at instruction following and JSON output
- Better than Qwen2.5-1.5B for structured medical extraction
- microsoft/Phi-3-mini-4k-instruct (~2.4GB download)
"""

import json, re, os
from functools import lru_cache
from dataclasses import dataclass
from typing import Optional
import torch

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
PHI3_MODEL_ID = "microsoft/Phi-3-mini-4k-instruct"

CORRECTIONS = {
    "imatimbing": "Imatinib", "imatims": "Imatinib", "imatimb": "Imatinib",
    "anahumb": "Imatinib",    "marquelo": "Imatinib",
    "riaston": "BCR-ABL",     "abilbits": "BCR-ABL", "bee-abl": "BCR-ABL",
}

EXTRACTION_PROMPT = """You are a medical records expert for Indian cancer hospitals.
Extract ALL information from this handwritten prescription OCR text.

CRITICAL: If OCR contains multiple DATE entries (28/11, 5/11/18, 24/12/18 etc),
each date is a SEPARATE VISIT. Create one record per date.

Abbreviations: T./Tab.=Tablet, c/o=complaints, O/E=examination, R/A=review after,
NAD=normal, OD=once daily, BD=twice, TDS=thrice, CML-CP=cancer type,
MMR=molecular response, BCR-ABL=blood test, CBC=blood count,
ECOG/PS=performance status, CNS/CVS/RS/PA=systems exam, Adv=advice.

Return ONLY valid JSON, nothing else:
{"records":[{"visit_date":"date or null","fields":{"field_name":"value"},"medications":[{"drug_name":"","dosage":"","frequency":"","route":""}]}],"document_type":"type","hospital_name":"name or null","patient_name":"name or null","hospital_no":"number or null"}"""


@dataclass
class ExtractionResult:
    success: bool
    records: list
    meta: dict
    mode: str = "text"
    error: Optional[str] = None


@lru_cache(maxsize=1)
def _load_phi3():
    """Load Phi-3 Mini once and cache. ~2.4GB download on first run."""
    from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
    tokenizer = AutoTokenizer.from_pretrained(PHI3_MODEL_ID, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        PHI3_MODEL_ID,
        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
    )
    model.to(DEVICE)
    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer,
                    device=DEVICE, trust_remote_code=True)
    return pipe


def _correct(text):
    low = text.lower()
    for wrong, right in CORRECTIONS.items():
        if wrong in low:
            text = re.sub(re.escape(wrong), right, text, flags=re.IGNORECASE)
    return text


def _correct_meds(meds):
    out = []
    for m in meds:
        if not isinstance(m, dict): continue
        m = dict(m)
        if m.get("drug_name"): m["drug_name"] = _correct(m["drug_name"])
        out.append(m)
    return out


def _extract_json(text):
    text = re.sub(r"```[a-z]*", "", text.strip()).replace("```", "").strip()
    start = text.find("{")
    if start == -1: return text
    depth = end = 0
    for i in range(start, len(text)):
        if text[i] == "{":   depth += 1
        elif text[i] == "}":
            depth -= 1
            if depth == 0: end = i; break
    if end: return text[start:end+1]
    partial = text[start:]
    ob  = partial.count("{") - partial.count("}")
    ob2 = partial.count("[") - partial.count("]")
    partial = re.sub(r',\s*"[^"]*"?\s*:?\s*"?[^"}\]]*$', "", partial)
    partial = re.sub(r',\s*$', "", partial)
    return partial + "]"*ob2 + "}"*ob


def _normalize(data):
    meta = {k: data.get(k) for k in
            ["document_type","hospital_name","patient_name","hospital_no"]}
    meta = {k: v for k,v in meta.items() if v}
    records = data.get("records", [])
    if not isinstance(records, list): records = [records]
    clean = []
    for rec in records:
        if not isinstance(rec, dict): continue
        fields = {k: str(v) for k,v in rec.get("fields",{}).items()
                  if v and str(v).strip() and k != "medications"}
        meds = rec.get("medications", [])
        if isinstance(meds, str):
            try: meds = json.loads(meds)
            except: meds = []
        if not isinstance(meds, list): meds = []
        meds = [m for m in meds if isinstance(m,dict) and
                any(v and str(v).strip() for v in m.values())]
        meds = _correct_meds(meds)
        clean.append({"visit_date": rec.get("visit_date"),
                      "fields": fields, "medications": meds})
    return clean, meta


def extract_fields(ocr_text: str, **kwargs) -> ExtractionResult:
    """Run Phi-3 Mini locally — no external API calls."""
    try:
        pipe = _load_phi3()
        messages = [
            {"role": "system", "content": EXTRACTION_PROMPT},
            {"role": "user",   "content": f"OCR TEXT:\n{ocr_text[:2000]}"},
        ]
        output = pipe(
            messages,
            max_new_tokens=1024,
            do_sample=False,
            temperature=None,
            top_p=None,
        )
        raw = output[0]["generated_text"]
        # Phi-3 returns full conversation — get last assistant message
        if isinstance(raw, list):
            raw = raw[-1].get("content", "")
        cleaned = _extract_json(str(raw))
        try:
            data = json.loads(cleaned)
        except json.JSONDecodeError:
            cleaned = re.sub(r'"\s*:\s*"([^"]*?)(?=[,}])', r'": "\1"', cleaned)
            data = json.loads(cleaned)
        records, meta = _normalize(data)
        return ExtractionResult(True, records, meta, mode="phi3-mini/local")
    except Exception as e:
        return ExtractionResult(False, [], {}, error=str(e))


def extract_from_image(image_bytes: bytes, **kwargs) -> ExtractionResult:
    return ExtractionResult(False, [], {}, mode="vision", error="NO_VISION_MODEL")