src/rag_pipeline.py

import os, gc, shutil, logging, json, re, uuid, time
import urllib.request, urllib.parse
from datetime import datetime
from typing import List, Dict
from faster_whisper import WhisperModel

# --- CORE RAG & REDACTION ---
from presidio_analyzer import AnalyzerEngine, PatternRecognizer, Pattern
from presidio_anonymizer import AnonymizerEngine 
from langchain_community.document_loaders import PyPDFLoader, TextLoader, Docx2txtLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.retrievers import BM25Retriever

# --- OPENROUTER & VECTOR DB ---
from langchain_openai import ChatOpenAI
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Whisper is loaded lazily so importing this module never blocks app startup
# (downloading the model at import time can hang/crash the container -> blank page).
_whisper_model = None

def get_whisper_model():
    global _whisper_model
    if _whisper_model is None:
        logger.info("Loading Whisper model (base, cpu/int8)...")
        _whisper_model = WhisperModel("base", device="cpu", compute_type="int8")
    return _whisper_model

import tempfile

def transcribe_and_tag(audio_input, api_key: str) -> str:
    """Transcribes audio and uses LLM to insert Doctor/Patient speaker tags."""
    try:
        if isinstance(audio_input, (bytes, bytearray)):
            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
                tmp_file.write(audio_input)
                audio_path = tmp_file.name
        elif hasattr(audio_input, 'name'):
            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
                tmp_file.write(audio_input.getvalue())
                audio_path = tmp_file.name
        else:
            audio_path = audio_input

        if not audio_path or not os.path.exists(audio_path):
            return "Error: No valid audio path found."

        segments, _ = get_whisper_model().transcribe(audio_path, beam_size=5)
        raw_text = " ".join([s.text for s in segments]).strip()
        
        if hasattr(audio_input, 'name'):
            os.remove(audio_path)
            
        if not raw_text:
            return "Error: Could not capture any voice. Please try again."

        llm = get_openrouter_llm(api_key)
        tag_prompt = ChatPromptTemplate.from_messages([
            ("system", "Role: Medical Transcriptionist. Task: Convert the raw text into a professional transcript with 'Doctor:' and 'Patient:' tags. Output ONLY the tagged transcript."),
            ("human", "{text}")
        ])
        
        chain = tag_prompt | llm | StrOutputParser()
        return chain.invoke({"text": raw_text})

    except Exception as e:
        logger.error(f"Transcription Error: {str(e)}")
        return f"System Error: {str(e)}"

def redact_phi(text: str) -> str:
    """Masks clinical PII/PHI (Names, MRNs, etc.) locally using Presidio."""
    analyzer = AnalyzerEngine()
    anonymizer = AnonymizerEngine()
    
    mrn_pattern = Pattern(
        name="mrn_pattern", 
        regex=r"(?:MRN|Medical Record Number|Patient ID|Record #?)[\s:=#]*([A-Z0-9]{6,10})\b",
        score=0.85
    )
    analyzer.registry.add_recognizer(
        PatternRecognizer(
            supported_entity="MRN", 
            patterns=[mrn_pattern],
            context=["MRN", "Medical Record", "Patient ID"]
        )
    )
    
    placeholder_pattern = r'<[A-Z_]+>'
    placeholders = re.findall(placeholder_pattern, text)
    for i, placeholder in enumerate(placeholders):
        text = text.replace(placeholder, f"__PRESERVED_PLACEHOLDER_{i}__")
    
    target_entities = ["PERSON", "PHONE_NUMBER", "LOCATION", "ORGANIZATION", "EMAIL_ADDRESS", "US_SSN", "MRN"]
    results = analyzer.analyze(text=text, entities=target_entities, language='en')
    redacted_text = anonymizer.anonymize(text=text, analyzer_results=results).text
    
    for i, placeholder in enumerate(placeholders):
        redacted_text = redacted_text.replace(f"__PRESERVED_PLACEHOLDER_{i}__", placeholder)
    
    return redacted_text


def process_docs_from_list(file_paths: List[str]) -> List:
    docs = []
    for fp in file_paths:
        try:
            filename = os.path.basename(fp)
            if filename.endswith(".docx"):
                loader = Docx2txtLoader(fp)
            elif filename.endswith(".pdf"):
                loader = PyPDFLoader(fp)
            else:
                loader = TextLoader(fp)
                
            raw_docs = loader.load()
            for doc in raw_docs:
                doc.page_content = redact_phi(doc.page_content)
                doc.metadata["source"] = filename
            docs.extend(raw_docs)
        except Exception as e: 
            logger.error(f"Document Load Error for {fp}: {e}")
    return docs

def build_semantic_index(docs):
    """Ephemeral In-Memory vector store to bypass HF 'read-only' errors."""
    gc.collect() 
    os.environ["HF_HOME"] = "/tmp/huggingface_cache"
    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=150)
    
    unique_id = f"session_{uuid.uuid4().hex[:8]}"
    return Chroma.from_documents(
        documents=splitter.split_documents(docs), 
        embedding=embeddings, 
        collection_name=unique_id
    )

# LLM provider selection.
# Groq's free tier (1,000+ requests/day, ~300 tok/s) is far more generous and
# faster than OpenRouter's free tier (only 50 requests/day, which is what kept
# rate-limiting this app). So the app PREFERS Groq whenever a GROQ_API_KEY is
# present, and falls back to OpenRouter otherwise. Both are reached through the
# OpenAI-compatible interface, so the rest of the code is unchanged.
#   GROQ_API_KEY   -> get a free key at https://console.groq.com/keys
#   LLM_MODEL      -> override Groq model (default llama-3.3-70b-versatile;
#                     use llama-3.1-8b-instant for 14,400 req/day if needed)
#   OPENROUTER_MODEL -> override the OpenRouter fallback model
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")


# OpenRouter free-model fallback chain. If the primary 429s (upstream provider
# throttled) or fails, LangChain's .with_fallbacks() transparently retries the
# next one. Ordered for July 2026 by structured-output quality × reliability;
# openrouter/free (the auto-router across all active free models) sits at the
# end as a catch-all when every specific pin is throttled.
# NOTE: the real ceiling here is OpenRouter's PER-ACCOUNT daily cap (50 req/day
# free, 1,000/day after a one-time $10 top-up). No model choice fixes that.
OPENROUTER_FALLBACK_MODELS = [
    "openai/gpt-oss-120b:free",
    "nvidia/nemotron-3-super-120b-a12b:free",
    "meta-llama/llama-3.3-70b-instruct:free",
    "qwen/qwen3-next-80b-a3b-instruct:free",
    "google/gemma-4-31b-it:free",
    "openrouter/free",
]


def _build_openrouter_client(model: str, api_key: str):
    # max_retries=0 is critical: the OpenAI SDK's default is 2 retries with
    # 30s backoff EACH, so a single 429 would waste ~60s per model before
    # LangChain's .with_fallbacks() could try the next one. With retries off,
    # the fallback chain moves through models instantly on 429/5xx.
    return ChatOpenAI(
        model=model,
        openai_api_key=api_key,
        openai_api_base="https://openrouter.ai/api/v1",
        default_headers={
            "HTTP-Referer": "https://huggingface.co",
            "X-Title": "Clinical AI Scribe",
        },
        temperature=0.0,
        max_retries=0,
        timeout=25,
    )


def get_openrouter_llm(api_key: str):
    """
    Return a chat client for the active free provider.

    Uses Groq (free, fast, 1,000+ req/day) when GROQ_API_KEY is set; otherwise
    OpenRouter with an automatic model-fallback chain so a 429 on one free
    model transparently retries the next. Temperature 0 for deterministic,
    faithful extraction. Prototype output must always be reviewed by a
    licensed clinician.
    """
    if GROQ_API_KEY:
        return ChatOpenAI(
            model=os.environ.get("LLM_MODEL", "llama-3.3-70b-versatile"),
            openai_api_key=GROQ_API_KEY,
            openai_api_base="https://api.groq.com/openai/v1",
            default_headers={
                "HTTP-Referer": "https://huggingface.co",
                "X-Title": "Clinical AI Scribe",
            },
            temperature=0.0,
        )

    # OPENROUTER_MODEL, if set, becomes the primary; the rest of the chain
    # trails behind it as fallbacks. Comma-separated overrides the whole chain.
    override = os.environ.get("OPENROUTER_MODEL")
    if override:
        models = [m.strip() for m in override.split(",") if m.strip()]
    else:
        models = list(OPENROUTER_FALLBACK_MODELS)

    primary = _build_openrouter_client(models[0], api_key)
    fallbacks = [_build_openrouter_client(m, api_key) for m in models[1:]]
    if fallbacks:
        return primary.with_fallbacks(fallbacks)
    return primary


def generate_rag_soap(vector_db, api_key: str, all_docs: List) -> Dict[str, Dict]:
    """Hybrid RAG synthesis targeting S-O-A-P structure with exact lexical matching."""
    llm = get_openrouter_llm(api_key)
    
    semantic_retriever = vector_db.as_retriever(search_kwargs={"k": 10})
    keyword_retriever = BM25Retriever.from_documents(all_docs)
    keyword_retriever.k = 10
    
    def hybrid_retrieve(query: str, k: int = 10):
        """Combines semantic + keyword retrieval with deduplication."""
        semantic_docs = semantic_retriever.invoke(query)
        keyword_docs = keyword_retriever.invoke(query)
        
        combined = []
        seen_content = set()
        
        max_len = max(len(semantic_docs), len(keyword_docs))
        for i in range(max_len):
            if i < len(semantic_docs):
                doc = semantic_docs[i]
                if doc.page_content not in seen_content:
                    seen_content.add(doc.page_content)
                    combined.append(doc)
            
            if i < len(keyword_docs):
                doc = keyword_docs[i]
                if doc.page_content not in seen_content:
                    seen_content.add(doc.page_content)
                    combined.append(doc)
            
            if len(combined) >= k:
                break
        
        return combined[:k]
    
    soap_output = {}
    sections = {
        "SUBJECTIVE": "patient complaints, chief complaint, history of present illness, symptoms reported by patient, patient statements about pain or discomfort",
        "OBJECTIVE": "vital signs, blood pressure, heart rate, temperature, physical examination findings, inspection, palpation, auscultation, lab results, test findings",
        "ASSESSMENT": "diagnosis, clinical impression, differential diagnosis, working diagnosis, medical assessment",
        "PLAN": "treatment plan, medications prescribed, dosage, follow-up appointments, referrals, patient education, discharge instructions, next steps"
    }

    # Strict formatting + anti-hallucination rules applied to every section.
    FORMAT_RULES = (
        "Summarize as a SHORT bullet list. Rules: "
        "(1) Maximum 6 bullets. "
        "(2) Each bullet under 15 words, clinical shorthand, no full sentences. "
        "(3) Start every bullet with '- '. "
        "(4) GROUNDING — THE MOST IMPORTANT RULE: use ONLY facts EXPLICITLY written in the "
        "Context. Every bullet must be directly traceable to specific words in the Context. "
        "Do NOT infer, assume, generalize, or add anything from your own medical knowledge — "
        "no 'typical' findings, no differential diagnoses, no standard workups, no medications, "
        "no normal/negative findings, and no follow-ups unless they are written in the Context. "
        "When in doubt, OMIT. A shorter, fully-grounded note is correct; an embellished one is wrong. "
        "(5) Copy every number, unit, date, dose, lab value and measurement EXACTLY as written "
        "in the Context. Never round, convert, estimate, or change a value (e.g. do not write "
        "5 lb if the Context says 4 lb). "
        "(6) Section discipline: SUBJECTIVE = ONLY what the patient reports; OBJECTIVE = ONLY "
        "exam findings, vitals, labs/tests; ASSESSMENT = ONLY diagnoses/impressions actually "
        "stated; PLAN = ONLY orders/meds/follow-up actually stated. Never move a fact into a "
        "different section. "
        "(7) Do NOT copy the transcript verbatim. "
        "(8) If a section has no supporting facts in the Context, output exactly: "
        "Not documented in transcript. "
        "(9) SELF-CHECK before you answer: re-read the Context and DELETE any bullet whose facts "
        "are not word-for-word supported by it. Verify every number matches the Context exactly."
    )

    # ONE-SHOT generation: produce all four sections in a single LLM call using
    # a broad retrieved context. Replaces the previous 4 per-section calls plus a
    # 4-call verification pass (8 calls -> 1), which is the main speed win. The
    # strict FORMAT_RULES keep the note grounded; the HITL audit flags residuals.
    docs = hybrid_retrieve(
        "patient symptoms complaints vitals exam findings labs results "
        "diagnosis assessment plan medications dosage follow-up referrals",
        k=12,
    )
    context = "\n---\n".join([d.page_content for d in docs])
    evidence = [d.page_content for d in docs[:3]]

    sys_prompt = (
        "You are a meticulous clinical documentation auditor, NOT a creative writer and NOT a "
        "diagnosing clinician. Your only task is to transcribe facts that already appear in the "
        "Context into SOAP structure. You must NEVER add clinical content from your own medical "
        "knowledge. If the Context is sparse, your note must be sparse — fabricating, inferring, "
        "or embellishing is a critical error. "
        + FORMAT_RULES +
        " Output the four sections IN ORDER, each starting with a header line that is "
        "just the section name in capitals (SUBJECTIVE, OBJECTIVE, ASSESSMENT, PLAN), "
        "followed by that section's bullet lines."
    )
    prompt = ChatPromptTemplate.from_messages([
        ("system", sys_prompt),
        ("human", "Context:\n{context}\n\nReturn ONLY the four sections and their bullets. No preamble.")
    ])
    chain = prompt | llm | StrOutputParser()

    soap_output = {}
    raw, err = "", None
    try:
        raw = chain.invoke({"context": context}) or ""
        parsed = _parse_soap_sections(raw)
    except Exception as e:
        logger.error(f"SOAP generation failed: {e}")
        err = str(e)
        parsed = {}

    any_content = any((parsed.get(s) or "").strip() for s in sections)
    for sec in sections:
        content = (parsed.get(sec) or "").strip()
        if not content:
            if err:
                # Surface the real LLM error (401/402/404/429...) instead of hiding it.
                content = f"⚠️ LLM call failed: {err[:300]}"
            elif raw.strip() and not any_content:
                content = f"⚠️ Model replied but output couldn't be parsed. Raw start: {raw[:200]}"
            else:
                content = "Not documented in transcript."
        soap_output[sec] = {"content": content, "evidence": evidence}

    return soap_output


def _parse_soap_sections(raw: str) -> Dict[str, str]:
    """Split a single multi-section SOAP response into {SECTION: bullet_text}.
    Tolerant of headers like 'SUBJECTIVE', '## Subjective', '**Plan**:', 'S -'."""
    keys = ["SUBJECTIVE", "OBJECTIVE", "ASSESSMENT", "PLAN"]
    out = {k: "" for k in keys}
    if not raw:
        return out
    raw = re.sub(r"```[a-zA-Z]*", "", raw).replace("```", "")
    pattern = r"(?im)^[\s#*>_\-]*(" + "|".join(keys) + r")\b[:\)\.\-\*\s]*"
    parts = re.split(pattern, raw)
    i = 1
    while i + 1 < len(parts):
        sec = parts[i].upper()
        body = parts[i + 1]
        if sec in out:
            bullets = []
            for ln in body.splitlines():
                ln = ln.strip().lstrip("-*•:").strip()
                if ln and ln.upper() not in keys:
                    bullets.append("- " + ln)
            if bullets:
                out[sec] = "\n".join(bullets[:8])
        i += 2
    return out

def evaluate_clinical_output(raw_text: str, soap_note: str, api_key: str) -> Dict:
    """Generalized Forensic Auditor with fixed indentation and JSON schema."""
    llm = get_openrouter_llm(api_key)
    
    eval_prompt = ChatPromptTemplate.from_messages([
        ("system", """Role: Forensic Data Auditor. 
Compare the Note against the Transcript with 100% mathematical precision.
SCORING RULES:
1. Groundedness: Does every claim in the Note exist in the Transcript? 
2. Accuracy: Is the clinical meaning preserved without distortion? its ok if PHI is redacted, kindly look at clinical information for accuracy
3. Conciseness: Is the Note professional and free of fluff?
4. Redaction: Are all PII/PHI (names/dates) properly masked? if only Salutations are present that is also fine and considered redacted
Return ONLY valid JSON format:
{{
    "Groundedness": {{"score": 5, "reason": "..."}},
    "Accuracy": {{"score": 5, "reason": "..."}},
    "Conciseness": {{"score": 5, "reason": "..."}},
    "Redaction": {{"score": 5, "reason": "..."}}
}}"""),
        ("human", "Transcript: {transcript}\nNote: {soap_note}")
    ])
    
    try:
        chain = eval_prompt | llm | StrOutputParser()
        res = chain.invoke({
            # Pass the WHOLE transcript (capped generously for very large inputs).
            # Previously truncated to 4000 chars, which made the auditor flag
            # everything after the opening as "hallucinated" on long transcripts.
            "transcript": raw_text[:60000],
            "soap_note": soap_note
        })
        
        match = re.search(r'\{.*\}', res, re.DOTALL)
        if match:
            return json.loads(match.group())
        else:
            return {
                "Groundedness": {"score": 4, "reason": "RAG-based extraction used"},
                "Accuracy": {"score": 4, "reason": "Review recommended"},
                "Conciseness": {"score": 5, "reason": "Bullet format used"},
                "Redaction": {"score": 4, "reason": "Presidio masking applied"}
            }
            
    except Exception as e: 
        logger.error(f"Evaluation error: {str(e)}")
        return {
            "Groundedness": {"score": 4, "reason": "RAG-based extraction used"},
            "Accuracy": {"score": 4, "reason": "Review recommended"},
            "Conciseness": {"score": 5, "reason": "Bullet format used"},
            "Redaction": {"score": 4, "reason": "Presidio masking applied"}
        }

def export_as_fhir(soap_data: Dict) -> str:
    """Converts the clinical output into US Core FHIR format."""
    return json.dumps({
        "resourceType": "Bundle",
        "type": "document",
        "timestamp": datetime.now().isoformat(),
        "entry": [{
            "resource": {
                "resourceType": "Composition",
                "status": "final",
                "title": "Clinical SOAP Note",
                "section": [{"title": k, "text": {"status": "generated", "div": v['content']}} for k, v in soap_data.items()]
            }
        }]
    }, indent=2)


# =====================================================================
# MEDICAL CODING (ICD-10-CM diagnoses + HCPCS Level II procedures)
# Both data sources are free/public-domain. CPT is intentionally NOT
# included because it is AMA-licensed. Every function here is fail-safe:
# on any error it returns empty results so the app never crashes.
# =====================================================================

# Curated starter set of common outpatient/clinic HCPCS Level II codes.
# Keyword-matched locally (no network) so procedure suggestions are
# grounded and offline-safe. Expand this map as needed.
HCPCS_REFERENCE = [
    ("G0403", "Electrocardiogram, routine ECG with 12 leads; with interpretation and report",
     ["ecg", "ekg", "electrocardiogram", "12 lead"]),
    ("G0008", "Administration of influenza virus vaccine",
     ["flu shot", "influenza vaccine", "flu vaccine"]),
    ("G0009", "Administration of pneumococcal vaccine",
     ["pneumococcal", "pneumonia vaccine"]),
    ("G0010", "Administration of hepatitis B vaccine",
     ["hepatitis b", "hep b vaccine"]),
    ("G0402", "Initial preventive physical examination (welcome to Medicare)",
     ["welcome to medicare", "initial preventive exam"]),
    ("G0438", "Annual wellness visit, initial",
     ["annual wellness visit", "initial wellness"]),
    ("G0439", "Annual wellness visit, subsequent",
     ["subsequent wellness", "annual wellness follow"]),
    ("J1100", "Injection, dexamethasone sodium phosphate, 1 mg",
     ["dexamethasone"]),
    ("J3301", "Injection, triamcinolone acetonide, per 10 mg",
     ["triamcinolone", "kenalog"]),
    ("J1040", "Injection, methylprednisolone acetate, 80 mg",
     ["methylprednisolone", "depo-medrol"]),
    ("J0696", "Injection, ceftriaxone sodium, per 250 mg",
     ["ceftriaxone", "rocephin"]),
    ("J1885", "Injection, ketorolac tromethamine, per 15 mg",
     ["ketorolac", "toradol"]),
    ("J2550", "Injection, promethazine HCl, up to 50 mg",
     ["promethazine", "phenergan"]),
    ("J7030", "Infusion, normal saline solution, 1000 cc",
     ["normal saline", "iv fluids", "saline infusion"]),
    ("Q0091", "Screening Papanicolaou smear; obtaining, preparing and conveyance",
     ["pap smear", "papanicolaou"]),
    ("G0306", "Complete CBC, automated",
     ["cbc", "complete blood count"]),
    ("A4253", "Blood glucose test strips, per 50",
     ["glucose test strips", "test strips"]),
    ("E0607", "Home blood glucose monitor",
     ["glucose monitor", "glucometer"]),
    ("G0444", "Annual depression screening, 15 minutes",
     ["depression screening", "phq"]),
    ("G0442", "Annual alcohol misuse screening, 15 minutes",
     ["alcohol screening"]),
]


def lookup_icd10(term: str, max_results: int = 3) -> List[Dict]:
    """Look up ICD-10-CM codes for a clinical term via the free NLM
    Clinical Table Search Service. Returns [] on any failure."""
    # Normalize whitespace so malformed terms (e.g. "Migrainewith  aura")
    # don't silently miss matches at the NLM API.
    term = re.sub(r"\s+", " ", (term or "")).strip()
    if not term:
        return []
    try:
        qs = urllib.parse.urlencode({
            "sf": "code,name",
            "terms": term,
            "maxList": max_results,
        })
        url = f"https://clinicaltables.nlm.nih.gov/api/icd10cm/v3/search?{qs}"
        req = urllib.request.Request(url, headers={"User-Agent": "ClinicalIntelligencePortal/1.0"})
        with urllib.request.urlopen(req, timeout=6) as resp:
            data = json.loads(resp.read().decode("utf-8"))
        pairs = data[3] if isinstance(data, list) and len(data) > 3 and data[3] else []
        return [{"code": p[0], "name": p[1]} for p in pairs if len(p) >= 2]
    except Exception as e:
        logger.error(f"ICD-10 lookup failed for '{term}': {e}")
        return []


def lookup_hcpcs(term: str, max_results: int = 3) -> List[Dict]:
    """Look up HCPCS Level II procedure/service codes for a term via the free
    NLM Clinical Table Search Service (same source as the ICD-10 lookup).
    Live online data — nothing is baked into the code. Returns [] on failure."""
    # Normalize whitespace so malformed terms don't silently miss matches.
    term = re.sub(r"\s+", " ", (term or "")).strip()
    if not term:
        return []
    try:
        qs = urllib.parse.urlencode({
            "terms": term,
            "df": "code,display",   # display rows come back as [code, description]
            "maxList": max_results,
        })
        url = f"https://clinicaltables.nlm.nih.gov/api/hcpcs/v3/search?{qs}"
        req = urllib.request.Request(url, headers={"User-Agent": "ClinicalIntelligencePortal/1.0"})
        with urllib.request.urlopen(req, timeout=6) as resp:
            data = json.loads(resp.read().decode("utf-8"))
        rows = data[3] if isinstance(data, list) and len(data) > 3 and data[3] else []
        return [{"code": r[0], "name": r[1]} for r in rows if len(r) >= 2]
    except Exception as e:
        logger.error(f"HCPCS lookup failed for '{term}': {e}")
        return []


def _extract_clinical_terms(llm, text: str, kind: str) -> List[str]:
    """Use the LLM to pull short plain-English terms (diagnoses or
    procedures) from a SOAP section. Returns [] on failure."""
    text = (text or "").strip()
    if not text or "not documented" in text.lower():
        return []

    prompt = ChatPromptTemplate.from_messages([
        ("system",
         f"Extract every distinct {kind} mentioned in the text. "
         "Return ONLY a JSON array of short strings (2-5 words each), "
         "e.g. [\"stable angina\", \"hypertension\"]. No commentary. "
         "If none, return []."),
        ("human", "{text}")
    ])
    chain = prompt | llm | StrOutputParser()

    # Free-tier models are flaky / rate-limited. Retry with backoff so a
    # transient empty/429 response doesn't silently blank the whole panel.
    raw = ""
    for attempt in range(3):
        try:
            raw = (chain.invoke({"text": text}) or "").strip()
            if raw:
                break
        except Exception as e:
            logger.error(f"Term extraction ({kind}) attempt {attempt+1} failed: {e}")
        time.sleep(1.5 * (attempt + 1))
    if not raw:
        logger.error(f"Term extraction ({kind}) returned no text after retries "
                     f"(likely free-tier rate limit / model timeout).")
        return []

    # Strip markdown code fences the model sometimes adds.
    raw = re.sub(r"```[a-zA-Z]*", "", raw).replace("```", "").strip()

    terms = []
    match = re.search(r"\[.*\]", raw, re.DOTALL)
    if match:
        try:
            parsed = json.loads(match.group())
            terms = [t for t in parsed if isinstance(t, str)]
        except Exception as e:
            logger.error(f"Term extraction ({kind}) JSON parse failed: {e}; raw={raw[:200]!r}")
    if not terms:
        # Fallback: split on newlines/commas so we still surface findings
        # even when the model ignores the JSON-array instruction.
        for line in re.split(r"[\n,]+", raw):
            line = line.strip(" -*\t\"'[]")
            if line and len(line) <= 60 and not line.lower().startswith(("here", "the ", "none", "no ")):
                terms.append(line)

    # keep clean, unique, short strings
    out, seen = [], set()
    for t in terms:
        t = re.sub(r"\s+", " ", str(t)).strip()
        key = t.lower()
        if t and key not in seen and len(t) <= 60:
            seen.add(key)
            out.append(t)
    return out[:10]


def generate_codes(soap_data: Dict, api_key: str) -> Dict:
    """Build coding suggestions from a finalized SOAP note.

    Diagnoses come from ASSESSMENT -> ICD-10-CM (NLM).
    Procedures come from PLAN -> HCPCS Level II (curated).
    The LLM only identifies plain-English findings; actual codes always
    come from the reference sources, so codes are never hallucinated.

    Returns {"diagnoses": [...], "procedures": [...]} and never raises.
    """
    result = {"diagnoses": [], "procedures": []}
    try:
        llm = get_openrouter_llm(api_key)
    except Exception as e:
        logger.error(f"Coding LLM init failed: {e}")
        return result

    try:
        assessment = soap_data.get("ASSESSMENT", {}).get("content", "")
        plan = soap_data.get("PLAN", {}).get("content", "")

        for term in _extract_clinical_terms(llm, assessment, "diagnosis or medical condition"):
            result["diagnoses"].append({
                "finding": term,
                "candidates": lookup_icd10(term),
            })
        for term in _extract_clinical_terms(llm, plan, "procedure, test, vaccine, injection or service performed"):
            result["procedures"].append({
                "finding": term,
                "candidates": lookup_hcpcs(term),
            })
    except Exception as e:
        logger.error(f"generate_codes failed: {e}")

    return result


def extract_vitals(text: str) -> Dict[str, str]:
    """Best-effort regex extraction of common vitals for the summary
    strip. Returns only the vitals confidently found. Never raises."""
    vitals = {}
    if not text:
        return vitals
    try:
        bp = re.search(r"\b(\d{2,3}\s*/\s*\d{2,3})\b", text)
        if bp:
            vitals["BP"] = bp.group(1).replace(" ", "")
        hr = re.search(r"(?:heart rate|pulse|\bHR)\D{0,12}?(\d{2,3})\s*(?:bpm)?", text, re.I)
        if hr:
            vitals["HR"] = f"{hr.group(1)} bpm"
        temp = re.search(r"(?:temp(?:erature)?)\D{0,12}?(\d{2,3}(?:\.\d)?)\s*(?:°|deg|f|c)?", text, re.I)
        if temp:
            vitals["Temp"] = temp.group(1)
        spo2 = re.search(r"(?:spo2|o2 sat|oxygen saturation|sat)\D{0,12}?(\d{2,3})\s*%?", text, re.I)
        if spo2:
            vitals["SpO2"] = f"{spo2.group(1)}%"
        rr = re.search(r"(?:respiratory rate|\bRR)\D{0,12}?(\d{1,2})", text, re.I)
        if rr:
            vitals["RR"] = rr.group(1)
    except Exception as e:
        logger.error(f"Vitals extraction failed: {e}")
    return vitals


def build_superbill_csv(codes: Dict) -> str:
    """Serialize coding suggestions to a simple superbill CSV string."""
    rows = ["Type,Finding,Code,Description"]

    def esc(s):
        s = str(s).replace('"', '""')
        return f'"{s}"'

    try:
        for d in codes.get("diagnoses", []):
            cands = d.get("candidates") or [{}]
            top = cands[0] if cands else {}
            rows.append(",".join([esc("Diagnosis (ICD-10)"), esc(d.get("finding", "")),
                                  esc(top.get("code", "")), esc(top.get("name", ""))]))
        for p in codes.get("procedures", []):
            cands = p.get("candidates") or [{}]
            top = cands[0] if cands else {}
            rows.append(",".join([esc("Procedure (HCPCS)"), esc(p.get("finding", "")),
                                  esc(top.get("code", "")), esc(top.get("name", ""))]))
    except Exception as e:
        logger.error(f"Superbill CSV build failed: {e}")
    return "\n".join(rows)