"""Dataset loading: patient registry, visit sessions, questions and checkpoints.

Each patient directory holds `dialogues.jsonl` (one record per visit) and `questions.jsonl`
(one record per question, carrying its checkpoint visit). Sessions are sorted by
(date, visit_number) and records without turns are dropped. A visit is rendered as
`=== Visit v{n} · {date} · {department} ===` followed by `Speaker: text` lines; annotation
fields (`meta`, per-turn `covers`) never enter memory text. A question's checkpoint must
name an existing visit, and its checkpoint date must equal that visit's date.
"""
from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path


def read_jsonl(path) -> list[dict]:
    return [json.loads(l) for l in open(path, encoding="utf-8") if l.strip()]


def load_sessions(path) -> list[dict]:
    """Return non-empty sessions sorted by (date, visit_number), keeping only the fields memory may see."""
    rows = read_jsonl(path)
    rows = [r for r in rows if r.get("turns")]
    rows.sort(key=lambda r: (r["date"], r["visit_number"]))
    return [{"visit_number": r["visit_number"], "date": r["date"],
             "department": r["department"], "session_id": r.get("session_id"),
             "turns": [{"speaker": t["speaker"], "text": t["text"]} for t in r["turns"]]}
            for r in rows]


def render_dialogue(rec: dict) -> str:
    return "\n".join(f"{t['speaker']}: {t['text']}" for t in rec["turns"])


def render_visit(rec: dict) -> str:
    return f"=== Visit v{rec['visit_number']} · {rec['date']} · {rec['department']} ===\n" + render_dialogue(rec)


@dataclass
class Question:
    qid: int
    raw: dict
    ckpt_visit: int = -1
    ckpt_date: str = ""

    @property
    def axis(self) -> str:
        return str(self.raw.get("axis", ""))

    @property
    def answer_type(self) -> str:
        return str(self.raw.get("answer_type", ""))

    @property
    def question(self) -> str:
        return self.raw["question"]

    @property
    def grounding(self) -> str:
        return str(self.raw.get("grounding", "record"))

    @property
    def evidence_visits(self) -> list[int]:
        return list(self.raw.get("evidence_visits", []))


def load_questions(questions_file, sessions: list[dict]) -> list[Question]:
    """Load a patient's questions, each with its checkpoint visit and date."""
    records = read_jsonl(questions_file)
    date_by_visit = {s["visit_number"]: s["date"] for s in sessions}
    questions, problems = [], []
    for i, r in enumerate(records):
        ck = r.get("checkpoint_visit")
        if ck is None:
            problems.append(f"qid {r.get('qid')}: no checkpoint_visit")
            continue
        ck = int(ck)
        if ck not in date_by_visit:
            problems.append(f"qid {r.get('qid')}: checkpoint_visit v{ck} not in dialogues")
            continue
        stated = r.get("checkpoint_date")
        if stated and str(stated) != date_by_visit[ck]:
            problems.append(f"qid {r.get('qid')}: checkpoint_date {stated} != v{ck} date {date_by_visit[ck]}")
            continue
        questions.append(Question(qid=i, raw=r, ckpt_visit=ck,
                                  ckpt_date=str(stated or date_by_visit[ck])))
    if problems:
        raise SystemExit("questions with checkpoints failed validation:\n  " + "\n  ".join(problems))
    return questions


QUESTIONS_FILE = "questions.jsonl"


@dataclass
class Patient:
    """One patient: give a `dir` and let `resolve()` find the files inside it, or set the paths explicitly."""
    name: str
    patient_id: str
    dialogues: str = ""
    questions_dir: str = ""
    dir: str = ""
    questions_file: str = ""

    def resolve(self) -> bool:
        """Fill unset paths from `dir`; return whether the questions file exists."""
        d = Path(self.dir) if self.dir else None
        if d and not self.dialogues and (d / "dialogues.jsonl").exists():
            self.dialogues = str(d / "dialogues.jsonl")
        if d and not self.questions_dir:
            self.questions_dir = str(d)
        qfile = Path(self.questions_dir) / QUESTIONS_FILE if self.questions_dir else None
        if qfile is None or not qfile.exists():
            return False
        self.questions_file = str(qfile)
        return True

    def load(self):
        """Return (sessions, questions) for this patient."""
        found = self.resolve()
        if not self.dialogues:
            raise SystemExit(f"patient '{self.name}': no dialogues file found (dir={self.dir})")
        if not found:
            raise SystemExit(f"patient '{self.name}': no questions file in {self.questions_dir} "
                             f"(looked for {QUESTIONS_FILE})")
        sessions = load_sessions(self.dialogues)
        questions = load_questions(self.questions_file, sessions)
        if not questions:
            raise SystemExit(f"patient '{self.name}': no questions found in {self.questions_file}")
        return sessions, questions


def load_patients(config_path, skipped: list | None = None) -> dict[str, Patient]:
    """Build the patient registry from a YAML file.

    `patients` lists explicit entries; `patient_globs` lists globs of patient directories, each
    registered under its directory name. Directories without a questions file are left out and,
    when `skipped` is given, appended to it as (name, reason).
    """
    import glob as _glob
    import yaml
    cfg = yaml.safe_load(open(config_path, encoding="utf-8")) or {}
    out: dict[str, Patient] = {}
    for p in cfg.get("patients") or []:
        out[p["name"]] = Patient(
            name=p["name"], patient_id=str(p["patient_id"]), dialogues=p.get("dialogues", ""),
            questions_dir=p.get("questions_dir", ""), dir=p.get("dir", ""))
    for entry in cfg.get("patient_globs") or []:
        pattern, prefix = (entry, "") if isinstance(entry, str) else (entry["path"], entry.get("prefix", ""))
        for path in sorted(_glob.glob(str(pattern))):
            d = Path(path)
            if not d.is_dir():
                continue
            pid = d.name
            name = f"{prefix}{pid}"
            if name in out:
                continue
            pat = Patient(name=name, patient_id=pid, dir=str(d))
            if not pat.resolve():
                if skipped is not None:
                    skipped.append((name, f"no {QUESTIONS_FILE} in {d}"))
                continue
            out[name] = pat
    return out