#!/usr/bin/env python3
"""Stage F: two-agent role-play dialogue generation over the event chain.

Per visit, the doctor agent sees the visit's atoms (flagged NEW or CARRIED), up to two prior visits in the
same department and a checklist of atoms to convey; the patient agent sees its own complaints, prior
same-department visits and an optional cross-department recap. The agents alternate as separate LLM
calls sharing the transcript until every atom is conveyed or the turn limit is reached; short visits get a
minimum length and every dialogue ends with a closing turn. The prompts are DIALOGUE_RULES, DOCTOR_SYSTEM and
PATIENT_SYSTEM in pipeline.prompts.

  generate --all --workers N --out dialogues.jsonl
  count-empty        number of visits whose dialogue has no turns
"""
from __future__ import annotations

import argparse
import json
import os
import re
import threading
import unicodedata
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from pipeline import prompts
from pipeline.common import BASE, FIELDS, model
from pipeline.llm import ChatClient

DATA = BASE
MODEL = model("dialogue")
SHARED = prompts.DIALOGUE_RULES
DOC_SYS = SHARED + "\n\n" + prompts.DOCTOR_SYSTEM
PAT_SYS = SHARED + "\n\n" + prompts.PATIENT_SYSTEM

_STOP = {"the", "of", "a", "an", "and", "or", "to", "with", "for", "no", "on", "at", "in", "by",
         "is", "was", "are", "s", "p", "his", "her", "you", "your"}
_ALIASES = {
    "lbp": "low back pain", "oa": "osteoarthritis", "mmt": "muscle strength",
    "pt": "physical therapy", "ot": "occupational therapy", "st": "speech therapy", "nsaid": "painkiller",
    "vas": "pain score", "r/o": "rule out", "bil": "bilateral", "gerd": "reflux",
    "hivd": "disc herniation", "ar": "allergic rhinitis", "u/e": "upper limb", "ldl": "cholesterol",
    "ck": "muscle enzyme", "rom": "range of motion", "bw": "weight", "bp": "blood pressure",
}
SMALL_TALK_TOPICS = ("traffic", "waiting time", "the hospital being busy", "feeling mildly nervous",
                     "coming from work", "the weather", "being worried while waiting for results")
PATIENT_SMALL_TALK_TOPICS = SMALL_TALK_TOPICS + ("getting to the clinic",)
_NUMBER = re.compile(r"\d+(?:\.\d+)?")
_VITAL_NAMES = {"BP": ("bp", "blood pressure"), "pulse": ("pulse", "hr", "heart rate", "pr"),
                "weight_kg": ("bw", "weight", "wt"), "height_cm": ("height", "ht", "bh")}


def norm(s: str) -> str:
    s = str(s).lower()
    for k, v in _ALIASES.items():
        s = re.sub(rf"(?<![a-z]){re.escape(k)}(?![a-z])", f"{k} {v}", s)
    return re.sub(r"\s+", " ", s)


def toks(s: str) -> set:
    return {t for t in re.findall(r"[a-z0-9]+", norm(s)) if len(t) > 2 and t not in _STOP}


def write_jsonl(path, records):
    with open(path, "w", encoding="utf-8") as fh:
        for r in records:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    return len(records)


def drop_caseless_letters(s):
    """The text without letters that have no case (Unicode category Lo), with empty parentheses and repeated spaces
    removed."""
    s = "".join(ch for ch in str(s) if unicodedata.category(ch) != "Lo")
    s = re.sub(r"\(\s*\)", "", s)
    return re.sub(r"\s{2,}", " ", re.sub(r"\s+\)", ")", re.sub(r"\(\s+", "(", s))).strip()


recs, tags = {}, {}


def load_chain():
    """Read the patient's event chain and change tags."""
    recs.update({r["visit_number"]: r for r in (json.loads(l) for l in open(DATA / "event_chain.jsonl"))})
    tags.update({r["visit_number"]: r for r in (json.loads(l) for l in open(DATA / "tags.jsonl"))})


def atoms_of(rec):
    out = []
    for f in FIELDS:
        for i, a in enumerate(rec.get(f) or []):
            out.append({"atom_id": f"v{rec['visit_number']}:{f}:{i}", "field": f, "text": drop_caseless_letters(a)})
    for k, val in (rec.get("vitals") or {}).items():
        if in_doctor_notes(rec, k, val):
            continue
        out.append({"atom_id": f"v{rec['visit_number']}:vitals:{k}", "field": "vitals", "text": f"{k}={val}"})
    return [a for a in out if a["text"]]


def in_doctor_notes(rec, key, val):
    """Whether a doctor_notes atom of the visit already records this vital sign: it names the sign and gives the
    same reading."""
    names = _VITAL_NAMES.get(key, (key.lower(),))
    value = re.compile(rf"(?<![\d.]){re.escape(str(val))}(?![\d.])")
    for a in rec.get("doctor_notes") or []:
        t = str(a).lower()
        if value.search(t) and any(re.search(rf"(?<![a-z]){re.escape(n)}(?![a-z])", t) for n in names):
            return True
    return False


def flag_novelty(vn, atoms):
    """NEW vs CARRIED: an atom is CARRIED when an atom of the same field in a prior same-department visit shares
    at least 70% of its words and has every number the atom states; a number that is new makes it NEW."""
    prev = tags.get(vn, {}).get("prev_visits", []) or []
    prior = {}
    for pv in prev:
        for pa in atoms_of(recs[pv]) if pv in recs else []:
            prior.setdefault(pa["field"], []).append((toks(pa["text"]), set(_NUMBER.findall(pa["text"]))))
    for a in atoms:
        at, an = toks(a["text"]), set(_NUMBER.findall(a["text"]))
        a["novelty"] = "CARRIED" if at and any(len(at & pt) / len(at) >= 0.7 and an <= pn
                                               for pt, pn in prior.get(a["field"], [])) else "NEW"


def window_txt(vn):
    prev = tags.get(vn, {}).get("prev_visits", []) or []
    if not prev:
        return "  none: this is the patient's first visit to this department."
    out = []
    for tag, pv in zip(("most recent", "one before"), reversed(prev[-2:])):
        if pv in recs:
            out.append(f"  [{tag}: v{pv} {recs[pv]['date']}] " + " | ".join(a["text"] for a in atoms_of(recs[pv])))
    return "\n".join(out)


def doctor_blob(rec, atoms):
    department, date = drop_caseless_letters(rec["department"]), rec["date"]
    p = [f"You are a {department} physician. Today's date is {date}.",
         "\n<current_visit_facts>: today's full record, one fact per line with its novelty flag. Put the atom "
         "id in covers."]
    for a in atoms:
        fl = "NEW    " if a["novelty"] == "NEW" else "CARRIED"
        atom_id, field, text = a["atom_id"], a["field"], a["text"]
        p.append(f"  {fl} | {atom_id} [{field}] {text}")
    p.append("\n<recent_visits_here>: your previous visits in this department, the source for recaps.")
    p.append(window_txt(rec["visit_number"]))
    p.append("\n<must_convey_this_visit>: every atom id above must be conveyed by the end of the visit.")
    return "\n".join(p)


def patient_blob(rec, atoms):
    p = ["You are the patient. You do not know today's exam findings, results, diagnosis or plan until the doctor "
         "tells you.", "\n<your_reasons_today>: your own complaints for today, in your own words. You may bring "
         "these up."]
    reasons = [(a["atom_id"], a["text"]) for a in atoms if a["field"] == "patient_self_description"]
    for atom_id, text in reasons:
        p.append(f"  {atom_id} {text}")
    if not reasons:
        p.append("  none: a routine follow-up with no new complaint of your own")
    p.append("\n<your_recent_visits_here>: your previous visits in this department. You lived through these.")
    p.append(window_txt(rec["visit_number"]))
    return "\n".join(p)


def cross_dept_recap(rec):
    """The patient's most recent visit to another department, as the <cross_dept_recap> line of the patient's
    opening <turn_task>, or ''."""
    cd = tags.get(rec["visit_number"], {}).get("prev_cross_dept_visit")
    if not (cd and cd in recs):
        return ""
    other_department, other_date = drop_caseless_letters(recs[cd]["department"]), recs[cd]["date"]
    return (f"\n<cross_dept_recap>: you recently saw a different department, {other_department}, on {other_date}. "
            "You may briefly mention that you went there; you know none of its findings or results.")


def agent_turn(client, system, blob, transcript, steering):
    conversation = "\n".join(f"{x['speaker']}: {x['text']}" for x in transcript) or "(the visit is just beginning)"
    user = (f"{blob}\n\n=== Conversation so far ===\n{conversation}\n\n=== Your turn: one short spoken turn ===\n"
            f"{steering}")
    o = client.complete_json(system, user)
    if not isinstance(o, dict) or "text" not in o:
        return {"text": str(o)[:200], "covers": [], "type": "clinical", "end_ok": False}
    cov = o.get("covers") or []
    return {"text": drop_caseless_letters(o.get("text", "")),
            "covers": [str(c) for c in cov] if isinstance(cov, list) else [],
            "type": str(o.get("type", "clinical")), "end_ok": bool(o.get("end_ok", False))}


WRAP = "  <steering>: this is the final wrap-up turn; you may set end_ok=true."
NOWRAP = "  <steering>: this is not the wrap-up turn; set end_ok=false."


def generate_dialogue(vn, client):
    rec = recs[vn]
    atoms = atoms_of(rec)
    flag_novelty(vn, atoms)
    ids = {a["atom_id"] for a in atoms}
    a_by_id = {a["atom_id"]: a for a in atoms}
    n_new = sum(1 for a in atoms if a["novelty"] == "NEW")
    small_talk_budget = int(tags.get(vn, {}).get("noise_score", 0) or 0)
    turns_target = max(5, min(48, 2 + 2 * min(n_new, 6) + (len(atoms) - n_new + 3) // 4 + small_talk_budget))
    d_blob, p_blob = doctor_blob(rec, atoms), patient_blob(rec, atoms)
    transcript, covered = [], set()
    small_talk_left = [small_talk_budget]
    used_topics = set()

    def small_talk_hint(patient=False):
        """The small-talk sentence of a turn: the speaker's next topic not yet used in the visit, starting at a
        position set by the visit number, while the visit's budget lasts."""
        if small_talk_left[0] > 0:
            small_talk_left[0] -= 1
            topics = PATIENT_SMALL_TALK_TOPICS if patient else SMALL_TALK_TOPICS
            order = [topics[(vn + k) % len(topics)] for k in range(len(topics))]
            t = next((x for x in order if x not in used_topics), order[0])
            used_topics.add(t)
            return (f" You may add one brief bit of small talk about {t}, an allowed topic. Keep it natural, and do "
                    "not repeat a small-talk topic already used.")
        return " No small talk this turn."

    def add(speaker, turn):
        covered.update(c for c in turn["covers"] if c in ids)
        transcript.append({"speaker": speaker, "type": turn.get("type", "clinical"),
                           "text": turn["text"], "covers": [c for c in turn["covers"] if c in ids]})

    def uncovered():
        u = [a for a in atoms if a["atom_id"] not in covered]
        return sorted(u, key=lambda a: (a["novelty"] != "NEW", a["field"] != "doctor_extra_notes"))

    def reasons():
        return "; ".join(f"[{a['atom_id']}] {a['text']}" for a in uncovered()
                         if a["field"] == "patient_self_description")

    def react():
        """The patient's turn after a doctor's clinical turn."""
        t = ("<turn_task>: phase=react. React to what the doctor just said, with a lay reaction or a short "
             "question.")
        own = reasons()
        if own:
            t += f" If you have not said them yet, you may voice these reasons and list their ids in covers: {own}."
        add("Patient", agent_turn(client, PAT_SYS, p_blob, transcript, t + small_talk_hint(patient=True) + "\n"
                                  + NOWRAP))

    def task(batch, phase, wrap=False):
        facts = "; ".join(f"[{a['atom_id']}] {a['text']}" for a in batch)
        t = f"<turn_task>: phase={phase}."
        if facts:
            t += f" Convey only these facts in this turn, and list their ids in covers: {facts}."
        return t + small_talk_hint() + ("\n" + (WRAP if wrap else NOWRAP))

    own = reasons()
    add("Patient", agent_turn(client, PAT_SYS, p_blob, transcript,
        "<turn_task>: phase=intake. Open the visit: say in your own everyday words why you are here today."
        + (f" Voice your reasons for today and list their ids in covers: {own}." if own else "")
        + small_talk_hint(patient=True) + cross_dept_recap(rec) + "\n" + NOWRAP))
    add("Doctor", agent_turn(client, DOC_SYS, d_blob, transcript,
        task([a for a in uncovered() if a["field"] != "patient_self_description"][:1], "intake")))

    guard = 0
    while uncovered() and len(transcript) < turns_target + 8 and guard < 40:
        guard += 1
        u = uncovered()
        batch = u[:1] if u[0]["novelty"] == "NEW" else u[:3]
        add("Doctor", agent_turn(client, DOC_SYS, d_blob, transcript, task(batch, "clinical")))
        if uncovered():
            react()

    n_recap = 0
    while uncovered() and n_recap < 6:
        n_recap += 1
        add("Doctor", agent_turn(client, DOC_SYS, d_blob, transcript,
            task(uncovered()[:5], "recap")))

    while len(transcript) < 4:
        add("Doctor", agent_turn(client, DOC_SYS, d_blob, transcript,
            "<turn_task>: phase=plan. Briefly restate today's plan or next step in plain words, using only plans "
            "already in <current_visit_facts>. Do not add new instructions, tests or follow-up intervals."
            + small_talk_hint()
            + "\n" + NOWRAP))
    add("Doctor", agent_turn(client, DOC_SYS, d_blob, transcript,
        "<turn_task>: phase=closing. Wrap the visit up warmly in one short line, without new instructions, "
        "tests or follow-up intervals." + WRAP))

    return {"visit_number": vn, "date": rec["date"], "department": rec["department"], "session_id": f"v{vn}",
            "turns": transcript,
            "meta": {"tag": tags.get(vn, {}).get("tag"), "noise_score": small_talk_budget,
                     "turns_target": turns_target, "n_atoms": len(atoms), "n_new": n_new, "engine": MODEL,
                     "atom_index": {a["atom_id"]: {"text": a["text"], "field": a["field"], "novelty": a["novelty"]}
                                    for a in atoms},
                     "coverage": {"ledger": len(ids), "covered": len(covered),
                                  "uncovered": [a_by_id[i]["atom_id"] for i in ids - covered]}}}


def generate(args) -> None:
    """Generate the dialogues of the selected visits, resuming from any existing output."""
    load_chain()
    vns = sorted(recs) if args.all else [int(x) for x in args.visits.split(",") if x.strip()]
    outpath = Path(args.out) if os.path.isabs(args.out) else DATA / args.out
    done = {}
    if outpath.exists() and not args.fresh:
        done = {r["visit_number"]: r for r in (json.loads(l) for l in open(outpath) if l.strip()) if r.get("turns")}
    todo = [v for v in vns if v not in done]
    print(f"visits: {len(vns)} total | already done {len(done)} | to generate {len(todo)}")
    client = ChatClient(model=MODEL, reasoning_effort="low", max_completion_tokens=800)
    lock = threading.Lock()
    n_ok = [0]
    n_err = [0]
    fh = open(outpath, "a", encoding="utf-8")

    def one(vn):
        try:
            r = generate_dialogue(vn, client)
        except Exception as e:
            r = {"visit_number": vn, "error": str(e)[:150], "turns": []}
        with lock:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
            fh.flush()
            (n_ok if r.get("turns") else n_err)[0] += 1
            if (n_ok[0] + n_err[0]) % 25 == 0:
                print(f"  progress {n_ok[0]+n_err[0]}/{len(todo)}  (ok {n_ok[0]}, err {n_err[0]})", flush=True)
        return r

    with ThreadPoolExecutor(max_workers=args.workers) as ex:
        list(ex.map(one, todo))
    fh.close()
    allr = {r["visit_number"]: r for r in (json.loads(l) for l in open(outpath) if l.strip())}
    out = sorted(allr.values(), key=lambda r: r["visit_number"])
    write_jsonl(outpath, out)
    ok = [r for r in out if r.get("turns")]
    err = [r for r in out if not r.get("turns")]
    tc = [len(r["turns"]) for r in ok]
    cov = [r["meta"]["coverage"] for r in ok if "meta" in r]
    print(f"\ndialogues: {len(ok)} ok, {len(err)} error/empty  turns min/med/max = "
          f"{min(tc)}/{sorted(tc)[len(tc)//2]}/{max(tc)}")
    print(f"visits <5 turns: {sum(1 for t in tc if t < 5)}")
    print(f"atom coverage: {sum(c['covered'] for c in cov)}/{sum(c['ledger'] for c in cov)} "
          f"({100*sum(c['covered'] for c in cov)/max(1,sum(c['ledger'] for c in cov)):.1f}%)")
    if err:
        print("visits with errors (retried on the next run):", [r["visit_number"] for r in err][:30])
    print(f"wrote {args.out}")


def count_empty() -> None:
    """Print the number of visits whose dialogue has no turns."""
    print(sum(1 for l in open(BASE / "dialogues.jsonl") if l.strip() and not json.loads(l).get("turns")))


def main():
    ap = argparse.ArgumentParser(description="Stage F: dialogue generation.")
    sp = ap.add_subparsers(dest="cmd", required=True)
    g = sp.add_parser("generate")
    g.add_argument("--all", action="store_true")
    g.add_argument("--visits", default="")
    g.add_argument("--workers", type=int, default=12)
    g.add_argument("--out", default="dialogues.jsonl")
    g.add_argument("--fresh", action="store_true", help="ignore any existing output and regenerate all")
    sp.add_parser("count-empty")
    a = ap.parse_args()
    generate(a) if a.cmd == "generate" else count_empty()


if __name__ == "__main__":
    main()