"""Locations, identifiers, models, question counts, progress reporting and the term glossary shared by the
stages.

build_patient.sh sets MEMEBENCH_WORK (the patient's working directory), MEMEBENCH_PATIENT (the patient's
number in the visit table), the per-patient facts and the model names. The visit table (the synthetic
timeline) and the term glossary are not distributed (MEMEBENCH_VISIT_TABLE, MEMEBENCH_GLOSSARY).
"""
from __future__ import annotations

import io
import json
import os
import re
import time
from functools import lru_cache
from pathlib import Path

VISIT_TABLE = Path(os.environ.get("MEMEBENCH_VISIT_TABLE", "<VISIT_TABLE>"))
GLOSSARY = Path(os.environ.get("MEMEBENCH_GLOSSARY", "<TERM_GLOSSARY>"))
SEX_MAP = {"0": "female", "1": "male"}
DEFAULT_PDESC = "longitudinal multi-department history"
FIELDS = ["diagnosis", "patient_self_description", "tests_and_medication", "doctor_notes",
          "doctor_plans", "health_advice", "doctor_extra_notes"]

# Question counts: the minimum of record questions on the axes that get synthetic questions (3, 5, 6, 7), and the
# sizes asked of steps axis1/axis2.
PER_AXIS_MINIMUM = {3: 3, 5: 10, 6: 10, 7: 10}
AXIS1_TARGET = "8-10"
AXIS1_FORM_MIX = "5 form=range + 3 form=never_again"
AXIS2_TARGET = "4"


def _required(name: str) -> str:
    value = os.environ.get(name, "")
    if not value or value.startswith("<"):
        raise SystemExit(f"{name} is not set; run the stages through build_patient.sh")
    return value


BASE = Path(_required("MEMEBENCH_WORK"))


def patient_id() -> int:
    """The patient's number in the visit table."""
    return int(_required("MEMEBENCH_PATIENT"))


def model(role: str) -> str:
    """The model configured for a role: extraction, generation or dialogue."""
    return _required(f"MEMEBENCH_{role.upper()}_MODEL")


def no_memory_exempt_axes() -> set:
    """Axis numbers whose questions are kept regardless of the no-memory probe (MEMEBENCH_NO_MEMORY_EXEMPT_AXES)."""
    if "MEMEBENCH_NO_MEMORY_EXEMPT_AXES" not in os.environ:
        raise SystemExit("MEMEBENCH_NO_MEMORY_EXEMPT_AXES is not set; run the stages through build_patient.sh")
    raw = os.environ["MEMEBENCH_NO_MEMORY_EXEMPT_AXES"]
    return {int(x) for x in raw.split(",") if x.strip()}


def _sex_str() -> str:
    """The patient's sex from the visit table (0 = female, 1 = male), for the synthetic-atom prompts."""
    patient = patient_id()
    try:
        import pandas as pd
        df = pd.read_parquet(VISIT_TABLE, columns=["record_id", "sex"])
        v = str(df[df["record_id"] == patient]["sex"].dropna().iloc[0]).strip()
    except Exception as e:
        raise RuntimeError(f"cannot read the patient's sex from {VISIT_TABLE}: {e}") from e
    if v not in SEX_MAP:
        raise RuntimeError(f"sex value {v!r} of the patient in {VISIT_TABLE} is not one of {sorted(SEX_MAP)}")
    return SEX_MAP[v]


@lru_cache(maxsize=None)
def patient_line() -> str:
    """The one-line patient description given to the synthetic-atom prompts."""
    pdesc = os.environ.get("MEMEBENCH_PDESC", DEFAULT_PDESC)
    return f"Patient ({_sex_str()}): {pdesc}."


def norm_value(v) -> str:
    """A timeline value for comparison: lower case, parentheses and a leading 'history of' removed."""
    v = re.sub(r"\s*\([^)]*\)", "", str(v or "")).strip().lower()
    v = re.sub(r"^(?:history of|hx of|h/o|known)\s+", "", v)
    return re.sub(r"\s+", " ", v).strip(" .;,")


def other_department_segments(timelines, entity, department):
    """Segments of the timelines that other departments keep for the same entity (name compared ignoring case and
    spacing)."""
    key = re.sub(r"\s+", " ", str(entity or "").strip().lower())
    return [s for t in timelines if t.get("department") != department
            and re.sub(r"\s+", " ", str(t.get("entity") or "").strip().lower()) == key
            for s in t.get("segments") or []]


def has_text(q, key="question") -> bool:
    """Whether a model reply carries a non-empty string under `key`."""
    return isinstance(q, dict) and isinstance(q.get(key), str) and bool(q[key].strip())


def lists_options(question, options) -> bool:
    """Whether a question text already lists its options: every option label appears as an item marker, '<label>.'
    or '<label>)', at the start of a line or after a space, '/' or '('."""
    labels = [str(o.get("label")) for o in options or [] if isinstance(o, dict) and o.get("label")]
    return bool(labels) and all(re.search(rf"(?:^|[\s/(]){re.escape(x)}[.)]\s", str(question or "")) for x in labels)


def patient_context(chain) -> str:
    """The record's departments, 15 most frequent diagnoses and 20 most frequent tests and medications, for the
    synthetic-atom prompts."""
    from collections import Counter
    depts = sorted({r["department"] for r in chain})
    dx = Counter(str(a)[:60] for r in chain for a in r.get("diagnosis") or [])
    tm = Counter(str(a)[:60] for r in chain for a in r.get("tests_and_medication") or [])
    return ("Departments: " + "; ".join(depts) + "\n"
            "Most frequent diagnoses: " + "; ".join(d for d, _ in dx.most_common(15)) + "\n"
            "Most frequent tests and medications: " + "; ".join(t for t, _ in tm.most_common(20)))


# Progress reporting for long steps: flushed progress lines and an atomic JSON heartbeat file.
def _hms(s: float) -> str:
    s = int(max(0, s))
    if s < 60:
        return f"{s}s"
    if s < 3600:
        return f"{s//60}m{s%60:02d}s"
    return f"{s//3600}h{(s%3600)//60:02d}m"


class Progress:
    def __init__(self, label: str, total: int, out: str | None = None,
                 every: int = 10, secs: float = 15.0):
        self.label, self.total, self.out = label, total, Path(out) if out else None
        self.every, self.secs = every, secs
        self.n = self.ok = self.fail = 0
        self.t0 = self.last_t = time.time()
        self.last_n = 0
        self.done = False

    def __enter__(self):
        self._say(force=True)
        return self

    def __exit__(self, *exc):
        self.done = True
        self._say(force=True, final=True)
        return False

    def tick(self, ok: bool = True, n: int = 1):
        self.n += n
        if ok:
            self.ok += n
        else:
            self.fail += n
        now = time.time()
        if (self.n - self.last_n) >= self.every or (now - self.last_t) >= self.secs:
            self._say()

    def _state(self, final=False):
        el = time.time() - self.t0
        rate = self.n / el if el > 0 else 0
        left = (self.total - self.n) / rate if rate > 0 and not final else 0
        return {
            "label": self.label, "done": self.done,
            "n": self.n, "total": self.total,
            "pct": round(100 * self.n / self.total, 1) if self.total else 0,
            "ok": self.ok, "fail": self.fail,
            "elapsed_s": round(el, 1), "rate_per_s": round(rate, 2),
            "eta_s": round(left, 1) if not final else 0,
            "updated": time.strftime("%H:%M:%S"),
        }

    def _say(self, force: bool = False, final: bool = False):
        st = self._state(final)
        self.last_t, self.last_n = time.time(), self.n
        bar_w = 24
        filled = int(bar_w * st["pct"] / 100)
        bar = "#" * filled + "." * (bar_w - filled)
        tail = f", {st['fail']} FAILED" if st["fail"] else ""
        eta = "done" if final else f"eta {_hms(st['eta_s'])}"
        print(f"  [{bar}] {st['pct']:>5.1f}%  {st['n']:>4}/{st['total']}  "
              f"{st['rate_per_s']:.2f}/s  {_hms(st['elapsed_s'])} elapsed, {eta}{tail}",
              flush=True)
        if self.out:
            self.out.parent.mkdir(parents=True, exist_ok=True)
            tmp = self.out.with_suffix(".tmp")
            tmp.write_text(json.dumps(st, indent=1))
            os.replace(tmp, self.out)


# Terminology normalization (GLOSSARY: term <TAB> normalized form): a single pass, longest
# match first; each term becomes "normalized form (term)".
def load_glossary():
    g = {}
    for ln in io.open(GLOSSARY, encoding="utf-8"):
        ln = ln.rstrip("\n")
        if ln.strip():
            src, norm = ln.split("\t")
            g[src] = norm
    return g


def glossary_re(g):
    keys = sorted(g, key=len, reverse=True)
    return re.compile("|".join(re.escape(k) for k in keys))


def apply_glossary(text, g, pat):
    return pat.sub(lambda m: f"{g[m.group(0)]} ({m.group(0)})", text)