#!/usr/bin/env python3
"""Prepare the TenaOS task-tagged LoRA corpus.

This module turns the assembled `final_dataset` folder into the report-aligned
training artifacts used by the Unsloth trainer:

1. audit    - inspect source JSONL files and write dataset readiness stats
2. normalize - canonicalize accepted traces and write accepted/rejected splits
3. build-sft - convert canonical traces to Gemma chat-template conversations
4. evaluate - run corpus-level readiness checks over normalized/split data
5. package  - write adapter metadata and merge/export instructions

The script is deliberately dependency-light so dataset readiness can be checked
on the local box before deciding whether training happens locally or on Vertex.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
from collections import Counter, defaultdict
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import lru_cache
from pathlib import Path
from typing import Any, Callable, Iterable


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_FINAL_DATASET = ROOT / "final_dataset"
DEFAULT_WORK_DIR = ROOT / "lora_training" / "artifacts"
REPORT_PATH = Path("/var/www/TenaOS/docs/technical-report.pdf")

# Root of the TenaAgent service tree, used ONLY to read the exact prompt .txt
# files it ships (see TENA_AGENT_PROMPTS_DIR below) — never imported as a
# Python package, so this script stays dependency-light per the module
# docstring. Override with the TENA_AGENT_SERVICE_ROOT env var if the
# checkout lives elsewhere.
import os as _os

DEFAULT_TENA_AGENT_SERVICE_ROOT = Path(
    _os.environ.get("TENA_AGENT_SERVICE_ROOT", "/var/www/ClinicDx_Demo/TenaAgent/service")
)
TENA_AGENT_PROMPTS_DIR = DEFAULT_TENA_AGENT_SERVICE_ROOT / "tena_agent_service" / "prompts"

# ---------------------------------------------------------------------------
# Direct imports of production tool-schema constants + the exact ephemeral
# tool-schema system-message builder.
#
# Unlike the prompt .txt files above (plain text, read verbatim), tool
# schemas are Python literals scattered across several modules. Importing
# them directly -- instead of hand-copying -- guarantees byte-for-byte
# parity with what llm_client.py actually puts on the wire and eliminates
# copy/drift risk. Verified these modules import standalone without
# triggering any DB/network/FastAPI-app side effects at import time.
#
# Found via the "Training vs Production ReAct Loop" audit (2026-07-02): every
# tool-calling turn in production unconditionally gets a SECOND system
# message prepended ahead of the task system prompt --
# llm_client.py::_messages_with_text_tool_instructions runs whenever
# `tools is not None`, and every call site below always passes non-None
# tools. The corpus previously only emitted the single task system prompt,
# so training never saw this message at all. See _tool_schema_system_message
# and its use in _make_example.
import sys as _sys

if str(DEFAULT_TENA_AGENT_SERVICE_ROOT) not in _sys.path:
    _sys.path.insert(0, str(DEFAULT_TENA_AGENT_SERVICE_ROOT))

from tena_agent_service.llm_client import (  # noqa: E402
    _messages_with_text_tool_instructions as _prod_tool_schema_wrapper,
)
from tena_agent_service.tool_loop import (  # noqa: E402
    TOOL_SCHEMAS as _CDS_TOOL_SCHEMAS,
    _MIN_SEARCHES as _CDS_MIN_SEARCHES,
)
from tena_agent_service.material_loop import (  # noqa: E402
    MATERIAL_TOOL_SCHEMAS as _MATERIAL_TOOL_SCHEMAS,
)
from tena_agent_service.form_pipeline.research_phase import (  # noqa: E402
    RESEARCH_TOOLS as _FORM_RESEARCH_TOOLS,
)
from tena_agent_service.form_pipeline.ciel_resolution import (  # noqa: E402
    _RESOLVE_TOOLS as _FORM_RESOLVE_TOOLS_RAW,
)
from tena_agent_service.report_pipeline.planning_phase import (  # noqa: E402
    PLAN_TOOLS as _REPORT_PLAN_TOOLS,
)
from tena_agent_service.report_pipeline.ciel_resolution import (  # noqa: E402
    REPORT_OPENAI_TOOLS as _REPORT_RESOLVE_TOOLS_RAW,
)
from tena_agent_service.agent_prompts import (  # noqa: E402
    apply_tool_description_overlay as _apply_tool_overlay,
)

# form_pipeline/ciel_resolution.py and report_pipeline/ciel_resolution.py both
# overlay per-workflow description text onto the raw tool list before sending
# it to the model (apply_tool_description_overlay). Apply the same overlay
# here so the tool-schema system message text matches exactly.
_FORM_RESOLVE_TOOLS = _apply_tool_overlay("form_builder", _FORM_RESOLVE_TOOLS_RAW)
_REPORT_RESOLVE_TOOLS = _apply_tool_overlay("report_builder", _REPORT_RESOLVE_TOOLS_RAW)

# material_loop.py hardcodes its minimum-search-count threshold inline (no
# named module constant) -- see the two `if search_count < 4:` sites in
# material_loop.py. Kept here as a named constant, checked against the
# module source at pin-check time by the corpus readiness evaluator.
_MATERIAL_MIN_SEARCHES = 4


@lru_cache(maxsize=None)
def _tool_schema_system_message(tool_schemas_json: str) -> dict[str, str]:
    """The exact ephemeral system message llm_client.py prepends ahead of the
    task system prompt on every tool-calling turn (see module docstring
    above). `tool_schemas_json` is a json-dumped tool list so this stays
    hashable for lru_cache; callers pass the list and we dump it here.
    """
    tool_schemas = json.loads(tool_schemas_json)
    return _prod_tool_schema_wrapper([], tool_schemas)[0]


def _tool_schema_message_for(tool_schemas: list[dict[str, Any]]) -> dict[str, str]:
    return _tool_schema_system_message(json.dumps(tool_schemas, sort_keys=True))

TRACE_KINDS = {
    "form",
    "report",
    "scribe_text_english",
    "cds",
    "patient_education",
    "voice_scribe_audio",
    "scribe_text_amharic",
}

TASK_TAGS = {
    "form": "[form]",
    "report": "[report]",
    "scribe_text_english": "[scribe]",
    "voice_scribe_audio": "[scribe]",
    "scribe_text_amharic": "[scribe-am]",
    "cds": "[cds]",
    "patient_education": "[edu]",
}

CANONICAL_SCHEMA = {
    "form": "form_trace_v1",
    "report": "report_trace_v1",
    "scribe_text_english": "scribe_trace_v1",
    "voice_scribe_audio": "voice_scribe_trace_v1",
    "scribe_text_amharic": "scribe_amharic_trace_v1",
    "cds": "cds_trace_v1",
    "patient_education": "patient_material_trace_v1",
}

PATIENT_EDUCATION_SECTIONS = (
    "## What You Have",
    "## Why It Matters",
    "## What To Do",
    "## Your Medications",
    "## What to Avoid",
    "## Follow-Up Schedule",
    "## When To Seek Help",
)

CDS_SECTIONS = (
    "## Clinical Assessment",
    "## Evidence-Based Considerations",
    "## Suggested Actions",
    "## Safety Alerts",
    "## Key Points",
)

GENERIC_EDU_PHRASES = (
    "your doctor has diagnosed you with a health condition",
    "health information for you",
    "ask your doctor for specific guidance",
)

PHI_PATTERNS = (
    re.compile(r"\b\d{3}[-.]\d{2}[-.]\d{4}\b"),  # US SSN-like
    re.compile(r"\b(?:phone|tel|mobile)\s*[:#]\s*\+?[\d\s().-]{7,}", re.I),
    re.compile(r"\b(?:patient name|name)\s*[:#]\s*[A-Z][a-z]+(?:\s+[A-Z][a-z]+)+\b", re.I),
    re.compile(r"\b(?:DOB|MRN|SSN)\s*[:#]\s*\S+", re.I),
)


@dataclass(frozen=True)
class DatasetEntry:
    kind: str
    output: Path
    accepted_count: int
    raw_seen: int
    task: str


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def load_manifest(final_dataset: Path) -> list[DatasetEntry]:
    manifest_path = final_dataset / "final_dataset_summary.json"
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    entries: list[DatasetEntry] = []
    for item in manifest.get("datasets") or []:
        output = Path(str(item["output"]))
        if not output.is_absolute():
            output = final_dataset.parent / output
        entries.append(
            DatasetEntry(
                kind=str(item["kind"]),
                output=output,
                accepted_count=int(item.get("accepted_count") or 0),
                raw_seen=int(item.get("raw_seen") or 0),
                task=str(item.get("task") or ""),
            )
        )
    return entries


def iter_jsonl(path: Path) -> Iterable[tuple[int, dict[str, Any] | None, str | None]]:
    with path.open("r", encoding="utf-8") as handle:
        for line_no, line in enumerate(handle, 1):
            if not line.strip():
                continue
            try:
                value = json.loads(line)
            except json.JSONDecodeError as exc:
                yield line_no, None, str(exc)
                continue
            if not isinstance(value, dict):
                yield line_no, None, "record_not_object"
                continue
            yield line_no, value, None


def write_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2, ensure_ascii=False, sort_keys=True) + "\n", encoding="utf-8")


def write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> int:
    path.parent.mkdir(parents=True, exist_ok=True)
    count = 0
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            count += 1
            handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")
    return count


def audit(final_dataset: Path, out_dir: Path) -> dict[str, Any]:
    entries = load_manifest(final_dataset)
    datasets: list[dict[str, Any]] = []
    for entry in entries:
        stats: dict[str, Any] = {
            "kind": entry.kind,
            "path": str(entry.output),
            "manifest_accepted_count": entry.accepted_count,
            "manifest_raw_seen": entry.raw_seen,
            "manifest_task": entry.task,
            "is_trace_kind": entry.kind in TRACE_KINDS,
            "json_errors": 0,
            "record_count": 0,
            "schema_versions": Counter(),
            "task_tags": Counter(),
            "ready_for_training": Counter(),
            "status_counts": Counter(),
            "decision_counts": Counter(),
            "missing_task_tag": 0,
            "missing_quality_ready": 0,
            "char_count_min": None,
            "char_count_max": 0,
            "char_count_total": 0,
            "sample_trace_ids": [],
        }
        for line_no, record, error in iter_jsonl(entry.output):
            if error:
                stats["json_errors"] += 1
                continue
            assert record is not None
            stats["record_count"] += 1
            encoded_len = len(json.dumps(record, ensure_ascii=False))
            stats["char_count_total"] += encoded_len
            stats["char_count_max"] = max(int(stats["char_count_max"]), encoded_len)
            if stats["char_count_min"] is None:
                stats["char_count_min"] = encoded_len
            else:
                stats["char_count_min"] = min(int(stats["char_count_min"]), encoded_len)
            if len(stats["sample_trace_ids"]) < 5:
                trace_id = record.get("trace_id") or record.get("example_id") or record.get("source_example_id") or f"line_{line_no}"
                stats["sample_trace_ids"].append(str(trace_id))
            stats["schema_versions"].update([str(record.get("schema_version") or "missing")])
            tag = record.get("task_tag")
            if tag:
                stats["task_tags"].update([str(tag)])
            elif entry.kind in TRACE_KINDS:
                stats["missing_task_tag"] += 1
            quality = record.get("quality") if isinstance(record.get("quality"), dict) else {}
            if "ready_for_training" in quality:
                stats["ready_for_training"].update([str(bool(quality.get("ready_for_training"))).lower()])
            elif entry.kind in TRACE_KINDS:
                stats["missing_quality_ready"] += 1
            if record.get("status"):
                stats["status_counts"].update([str(record.get("status"))])
            decision = nested_get(record, ("comparison", "decision"))
            if decision:
                stats["decision_counts"].update([str(decision)])
        record_count = int(stats["record_count"])
        stats["char_count_avg"] = round(float(stats["char_count_total"]) / record_count, 2) if record_count else 0
        for key in ("schema_versions", "task_tags", "ready_for_training", "status_counts", "decision_counts"):
            stats[key] = dict(sorted(stats[key].items()))
        datasets.append(stats)
    result = {
        "schema_version": "tenaos_lora_dataset_audit_v1",
        "created_at": utc_now(),
        "technical_report": str(REPORT_PATH),
        "final_dataset": str(final_dataset),
        "trace_record_count": sum(item["record_count"] for item in datasets if item["is_trace_kind"]),
        "all_record_count": sum(item["record_count"] for item in datasets),
        "datasets": datasets,
        "report_alignment_notes": [
            "Only records with canonical task_tag and quality.ready_for_training should feed LoRA training.",
            "Auxiliary scribe prompt corpora are audited but excluded from trace SFT by default.",
            "The report target includes Amharic voice traces; the assembled final dataset does not expose a separate Amharic voice kind.",
        ],
    }
    write_json(out_dir / "audit" / "dataset_audit.json", result)
    return result


def nested_get(record: dict[str, Any], path: tuple[str, ...]) -> Any:
    value: Any = record
    for key in path:
        if not isinstance(value, dict):
            return None
        value = value.get(key)
    return value


def normalize(final_dataset: Path, out_dir: Path) -> dict[str, Any]:
    entries = [entry for entry in load_manifest(final_dataset) if entry.kind in TRACE_KINDS]
    normalizers: dict[str, Callable[[dict[str, Any], int], tuple[dict[str, Any] | None, list[str]]]] = {
        "form": normalize_form,
        "report": normalize_report,
        "scribe_text_english": normalize_scribe_text,
        "cds": normalize_cds,
        "patient_education": normalize_patient_education,
        "voice_scribe_audio": normalize_voice_scribe,
        "scribe_text_amharic": normalize_scribe_amharic,
    }
    summary: dict[str, Any] = {
        "schema_version": "tenaos_lora_normalization_summary_v1",
        "created_at": utc_now(),
        "source_final_dataset": str(final_dataset),
        "output_root": str(out_dir / "normalized"),
        "datasets": [],
    }
    all_accepted_count = 0
    for entry in entries:
        accepted: list[dict[str, Any]] = []
        rejected: list[dict[str, Any]] = []
        reason_counts: Counter[str] = Counter()
        for line_no, record, error in iter_jsonl(entry.output):
            if error:
                reasons = ["invalid_json", error]
                rejected.append(reject_record(entry.kind, line_no, {"line": ""}, reasons))
                reason_counts.update(reasons)
                continue
            assert record is not None
            canonical, reasons = normalizers[entry.kind](record, len(accepted) + 1)
            reasons.extend(common_training_reject_reasons(record))
            reasons = sorted(set(reasons))
            if reasons:
                rejected.append(reject_record(entry.kind, line_no, record, reasons))
                reason_counts.update(reasons)
            else:
                assert canonical is not None
                canonical["trace_id"] = f"{entry.kind}_{len(accepted) + 1:06d}"
                accepted.append(canonical)
        kind_dir = out_dir / "normalized" / entry.kind
        accepted_path = kind_dir / f"{entry.kind}_accepted.jsonl"
        rejected_path = kind_dir / f"{entry.kind}_rejected.jsonl"
        write_jsonl(accepted_path, accepted)
        write_jsonl(rejected_path, rejected)
        all_accepted_count += len(accepted)
        item = {
            "kind": entry.kind,
            "task_tag": TASK_TAGS[entry.kind],
            "canonical_schema_version": CANONICAL_SCHEMA[entry.kind],
            "source_path": str(entry.output),
            "accepted_path": str(accepted_path),
            "rejected_path": str(rejected_path),
            "source_records": len(accepted) + len(rejected),
            "accepted_count": len(accepted),
            "rejected_count": len(rejected),
            "accept_rate": round(len(accepted) / (len(accepted) + len(rejected)), 4) if accepted or rejected else 0,
            "reject_reasons": dict(sorted(reason_counts.items())),
        }
        summary["datasets"].append(item)
        write_json(kind_dir / f"{entry.kind}_summary.json", item)
    summary["total_accepted_records"] = all_accepted_count
    write_json(out_dir / "normalized" / "normalization_summary.json", summary)
    return summary


def reject_record(kind: str, line_no: int, record: dict[str, Any], reasons: list[str]) -> dict[str, Any]:
    return {
        "schema_version": "tenaos_lora_reject_v1",
        "kind": kind,
        "line_no": line_no,
        "reject_reasons": sorted(set(reasons)),
        "raw": record,
    }


def common_training_reject_reasons(record: dict[str, Any]) -> list[str]:
    reasons: list[str] = []
    text = "\n".join(extract_narrative_text(record))
    if any(pattern.search(text) for pattern in PHI_PATTERNS):
        reasons.append("phi_like_text")
    for item in iter_extracted_items(record):
        if not isinstance(item, dict):
            continue
        if item.get("retired") is True:
            reasons.append("retired_ciel_concept")
        if str(item.get("resolutionStatus") or "").lower() == "unresolved":
            reasons.append("unresolved_item")
    return reasons


def extract_narrative_text(value: Any, *, parent_key: str = "") -> list[str]:
    """Return free-text fields where PHI checks are meaningful.

    Full traces contain concept IDs, timestamps, UUIDs, and serialized tool
    payloads; scanning every byte creates false positives. Restrict PHI scans to
    user/source narratives and generated clinical prose.
    """
    narrative_keys = {
        "request",
        "description",
        "noteText",
        "patientSummary",
        "content",
        "soap_text",
        "subjective",
        "objective",
        "assessment",
        "plan",
    }
    if isinstance(value, dict):
        out: list[str] = []
        for key, child in value.items():
            if key in narrative_keys and isinstance(child, str):
                out.append(child)
            elif isinstance(child, (dict, list)):
                out.extend(extract_narrative_text(child, parent_key=key))
        return out
    if isinstance(value, list):
        out = []
        for child in value:
            if isinstance(child, (dict, list)):
                out.extend(extract_narrative_text(child, parent_key=parent_key))
        return out
    return []


def iter_extracted_items(record: dict[str, Any]) -> Iterable[Any]:
    containers = [
        record.get("result"),
        nested_get(record, ("text_trace", "result")),
        nested_get(record, ("audio_trace", "result")),
        nested_get(record, ("english_trace", "result")),
        nested_get(record, ("amharic_trace", "result")),
    ]
    for container in containers:
        if not isinstance(container, dict):
            continue
        for key in ("concepts", "observations", "medications"):
            values = container.get(key)
            if isinstance(values, list):
                yield from values


def normalize_form(raw: dict[str, Any], index: int) -> tuple[dict[str, Any] | None, list[str]]:
    reasons: list[str] = []
    if raw.get("schema_version") != "form_trace_v1":
        reasons.append("invalid_schema_version")
    if raw.get("task_tag") != "[form]":
        reasons.append("missing_or_wrong_task_tag")
    quality = raw.get("quality") if isinstance(raw.get("quality"), dict) else {}
    if quality.get("ready_for_training") is not True:
        reasons.append("not_ready_for_training")
    if reasons:
        return None, reasons
    record = dict(raw)
    record["task"] = "form"
    record["task_tag"] = "[form]"
    record["schema_version"] = "form_trace_v1"
    record["quality"] = {**quality, "ready_for_training": True, "canonicalizer": "prepare_lora_corpus_v1"}
    return record, []


def normalize_report(raw: dict[str, Any], index: int) -> tuple[dict[str, Any] | None, list[str]]:
    reasons: list[str] = []
    if raw.get("schema_version") != "production_report_pipeline_trace_v1":
        reasons.append("invalid_schema_version")
    if raw.get("status") != "completed":
        reasons.append("pipeline_failed")
    if raw.get("error"):
        reasons.append("generation_error")
    request = raw.get("source_request") if isinstance(raw.get("source_request"), dict) else {}
    draft = raw.get("draft") if isinstance(raw.get("draft"), dict) else {}
    summary = raw.get("summary") if isinstance(raw.get("summary"), dict) else {}
    events = raw.get("events") if isinstance(raw.get("events"), list) else []
    spec = draft.get("spec") if isinstance(draft.get("spec"), dict) else {}
    if not request.get("example_id"):
        reasons.append("missing_source_example_id")
    if not summary.get("compiled"):
        reasons.append("not_compiled")
    if not summary.get("ran_report"):
        reasons.append("report_not_run")
    if len(spec.get("filters") or []) < 1:
        reasons.append("missing_filters")
    if len(events) < 4:
        reasons.append("too_few_events")
    if draft.get("conversationState") == "awaiting_question" and not draft.get("lastQuery"):
        reasons.append("incomplete_report_draft")
    if reasons:
        return None, reasons
    return canonical_record(
        "report",
        source={"type": "production_pipeline", "model": "gemma_4_e4b"},
        payload={
            "source_request": request,
            "request_metadata": request.get("metadata") if isinstance(request.get("metadata"), dict) else {},
            "draft": draft,
            "events": events,
            "summary": summary,
        },
        quality={"compiled": True, "ran_report": True, "validator": "prepare_lora_report_gate_v1"},
    ), []


def normalize_cds(raw: dict[str, Any], index: int) -> tuple[dict[str, Any] | None, list[str]]:
    reasons: list[str] = []
    if raw.get("schema_version") not in {"vertex_cds_trace_v1", "production_cds_trace_v1"}:
        reasons.append("invalid_schema_version")
    if raw.get("status") != "completed":
        reasons.append("pipeline_failed")
    if raw.get("error"):
        reasons.append("generation_error")
    source_case = raw.get("source_case") if isinstance(raw.get("source_case"), dict) else {}
    trace = raw.get("trace") if isinstance(raw.get("trace"), dict) else {}
    structured = raw.get("structured_cds") if isinstance(raw.get("structured_cds"), dict) else {}
    summary = raw.get("summary") if isinstance(raw.get("summary"), dict) else {}
    content = str(structured.get("content") or "")
    if not source_case.get("case_id"):
        reasons.append("missing_case_id")
    if len(trace.get("events") or []) < 4:
        reasons.append("too_few_events")
    if int(summary.get("search_count") or 0) < 4:
        reasons.append("too_few_searches")
    if any(section not in content for section in CDS_SECTIONS):
        reasons.append("missing_cds_section")
    if len(content) < 1200:
        reasons.append("content_too_short")
    if reasons:
        return None, reasons
    return canonical_record(
        "cds",
        source={"type": "vertex_production_pipeline", "model": "gemma_4_e4b"},
        payload={
            "source_case": source_case,
            "case_metadata": source_case.get("metadata") if isinstance(source_case.get("metadata"), dict) else {},
            "trace": trace,
            "structured_cds": structured,
            "summary": summary,
        },
        quality={"validator": "prepare_lora_cds_gate_v1"},
    ), []


def normalize_patient_education(raw: dict[str, Any], index: int) -> tuple[dict[str, Any] | None, list[str]]:
    reasons: list[str] = []
    if raw.get("schema_version") not in {"vertex_patient_material_trace_v1", "production_patient_material_trace_v1"}:
        reasons.append("invalid_schema_version")
    if raw.get("status") != "completed":
        reasons.append("pipeline_failed")
    if raw.get("error"):
        reasons.append("generation_error")
    source_case = raw.get("source_case") if isinstance(raw.get("source_case"), dict) else {}
    trace = raw.get("trace") if isinstance(raw.get("trace"), dict) else {}
    material = raw.get("material") if isinstance(raw.get("material"), dict) else {}
    summary = raw.get("summary") if isinstance(raw.get("summary"), dict) else {}
    content = str(material.get("content") or "")
    if not source_case.get("case_id"):
        reasons.append("missing_case_id")
    if len(trace.get("events") or []) < 4:
        reasons.append("too_few_events")
    if int(summary.get("search_count") or 0) < 4:
        reasons.append("too_few_searches")
    if any(section not in content for section in PATIENT_EDUCATION_SECTIONS):
        reasons.append("missing_material_section")
    if len(content) < 1600:
        reasons.append("content_too_short")
    lowered = content.lower()
    if any(phrase in lowered for phrase in GENERIC_EDU_PHRASES):
        reasons.append("generic_or_fallback_material")
    if reasons:
        return None, reasons
    return canonical_record(
        "patient_education",
        source={"type": "vertex_production_pipeline", "model": "gemma_4_e4b"},
        payload={
            "source_case": source_case,
            "case_metadata": source_case.get("metadata") if isinstance(source_case.get("metadata"), dict) else {},
            "trace": trace,
            "material": material,
            "summary": summary,
        },
        quality={"validator": "prepare_lora_patient_education_gate_v1"},
    ), []


def normalize_scribe_text(raw: dict[str, Any], index: int) -> tuple[dict[str, Any] | None, list[str]]:
    reasons = validate_scribe_trace(raw, schema_versions={"vertex_scribe_trace_v1", "production_scribe_trace_v1"})
    if reasons:
        return None, reasons
    sample = raw.get("source_sample") if isinstance(raw.get("source_sample"), dict) else {}
    return canonical_record(
        "scribe_text_english",
        source={"type": "vertex_production_pipeline", "model": "gemma_4_e4b"},
        payload={
            "source_sample": sample,
            "sample_metadata": sample.get("metadata") if isinstance(sample.get("metadata"), dict) else {},
            "result": raw.get("result"),
            "soap_text": raw.get("soap_text") or "",
            "summary": raw.get("summary") if isinstance(raw.get("summary"), dict) else {},
        },
        quality={"validator": "prepare_lora_scribe_text_gate_v1"},
    ), []


def validate_scribe_trace(raw: dict[str, Any], *, schema_versions: set[str]) -> list[str]:
    reasons: list[str] = []
    if raw.get("schema_version") not in schema_versions:
        reasons.append("invalid_schema_version")
    if raw.get("status") != "completed":
        reasons.append("pipeline_failed")
    if raw.get("error"):
        reasons.append("generation_error")
    sample = raw.get("source_sample") if isinstance(raw.get("source_sample"), dict) else {}
    result = raw.get("result") if isinstance(raw.get("result"), dict) else {}
    summary = raw.get("summary") if isinstance(raw.get("summary"), dict) else {}
    soap = result.get("soap") if isinstance(result.get("soap"), dict) else {}
    trace = result.get("generationTrace") if isinstance(result.get("generationTrace"), list) else []
    if not sample.get("example_id"):
        reasons.append("missing_source_example_id")
    for key in ("subjective", "objective", "assessment", "plan"):
        value = str(soap.get(key) or "").strip()
        if not value or value.lower() == "not documented":
            reasons.append(f"missing_soap_{key}")
    if len(trace) < 2:
        reasons.append("too_few_trace_events")
    if int(summary.get("unresolved_items") or 0) > 0:
        reasons.append("unresolved_items")
    if not any(result.get(key) for key in ("concepts", "observations", "medications")):
        reasons.append("empty_extraction")
    return reasons


def normalize_voice_scribe(raw: dict[str, Any], index: int) -> tuple[dict[str, Any] | None, list[str]]:
    reasons: list[str] = []
    if raw.get("schema_version") != "paired_text_audio_scribe_trace_v1":
        reasons.append("invalid_schema_version")
    if nested_get(raw, ("comparison", "decision")) != "accepted":
        reasons.append("paired_gate_not_accepted")
    audio_trace = raw.get("audio_trace") if isinstance(raw.get("audio_trace"), dict) else {}
    text_trace = raw.get("text_trace") if isinstance(raw.get("text_trace"), dict) else {}
    source_audio = raw.get("source_audio") if isinstance(raw.get("source_audio"), dict) else {}
    missing = nested_get(raw, ("comparison", "missing_from_audio")) or {}
    if any(missing.get(key) for key in ("concepts", "observations", "medications")):
        reasons.append("audio_missing_text_items")
    if audio_trace.get("status") != "completed" or text_trace.get("status") != "completed":
        reasons.append("paired_trace_failed")
    if int(nested_get(audio_trace, ("summary", "unresolved_items")) or 0) > 0:
        reasons.append("audio_unresolved_items")
    if int(source_audio.get("sample_rate") or 0) != 16000:
        reasons.append("invalid_sample_rate")
    if int(source_audio.get("channels") or 0) != 1:
        reasons.append("invalid_channels")
    audio_path = str(source_audio.get("audio_path") or "")
    if not audio_path:
        reasons.append("missing_audio_path")
    elif not resolve_audio_path(audio_path).exists():
        reasons.append("audio_file_missing")
    if reasons:
        return None, reasons
    return canonical_record(
        "voice_scribe_audio",
        source={"type": "paired_text_audio_gate", "model": "gemma_4_e4b_vertex", "tts": source_audio.get("voice")},
        payload={
            "source_audio": source_audio,
            "audio_metadata": source_audio.get("metadata") if isinstance(source_audio.get("metadata"), dict) else {},
            "text_trace": text_trace,
            "audio_trace": audio_trace,
            "comparison": raw.get("comparison"),
            "summary": audio_trace.get("summary") if isinstance(audio_trace, dict) else {},
        },
        quality={"validator": "prepare_lora_voice_scribe_gate_v1"},
    ), []


def resolve_audio_path(path: str) -> Path:
    candidate = Path(path)
    if candidate.is_absolute():
        return candidate
    return ROOT / candidate


def normalize_scribe_amharic(raw: dict[str, Any], index: int) -> tuple[dict[str, Any] | None, list[str]]:
    reasons: list[str] = []
    if raw.get("schema_version") != "paired_english_amharic_scribe_trace_v1":
        reasons.append("invalid_schema_version")
    if nested_get(raw, ("comparison", "decision")) != "accepted":
        reasons.append("paired_gate_not_accepted")
    english = raw.get("english_trace") if isinstance(raw.get("english_trace"), dict) else {}
    amharic = raw.get("amharic_trace") if isinstance(raw.get("amharic_trace"), dict) else {}
    if english.get("status") != "completed" or amharic.get("status") != "completed":
        reasons.append("paired_trace_failed")
    if int(nested_get(english, ("summary", "unresolved_items")) or 0) > 0:
        reasons.append("english_reference_unresolved")
    if int(nested_get(amharic, ("summary", "unresolved_items")) or 0) > 0:
        reasons.append("amharic_unresolved_items")
    missing = nested_get(raw, ("comparison", "missing_from_amharic")) or {}
    if any(missing.get(key) for key in ("concepts", "observations", "medications")):
        reasons.append("amharic_missing_english_items")
    if reasons:
        return None, reasons
    source_pair = raw.get("source_pair") if isinstance(raw.get("source_pair"), dict) else {}
    return canonical_record(
        "scribe_text_amharic",
        source={"type": "paired_english_amharic_gate", "model": "gemma_4_e4b_vertex"},
        payload={
            "source_pair": source_pair,
            "pair_metadata": source_pair.get("metadata") if isinstance(source_pair.get("metadata"), dict) else {},
            "english_trace": english,
            "amharic_trace": amharic,
            "comparison": raw.get("comparison"),
            "summary": amharic.get("summary") if isinstance(amharic, dict) else {},
        },
        quality={"validator": "prepare_lora_scribe_amharic_gate_v1"},
    ), []


def canonical_record(kind: str, *, source: dict[str, Any], payload: dict[str, Any], quality: dict[str, Any]) -> dict[str, Any]:
    return {
        "schema_version": CANONICAL_SCHEMA[kind],
        "trace_id": "",
        "task": canonical_task(kind),
        "task_tag": TASK_TAGS[kind],
        "source": source,
        **payload,
        "quality": {
            **quality,
            "ready_for_training": True,
            "canonicalizer": "prepare_lora_corpus_v1",
        },
    }


def canonical_task(kind: str) -> str:
    return {
        "form": "form",
        "report": "report",
        "scribe_text_english": "scribe",
        "voice_scribe_audio": "voice_scribe",
        "scribe_text_amharic": "scribe_amharic",
        "cds": "cds",
        "patient_education": "patient_education",
    }[kind]


# ---------------------------------------------------------------------------
# Multi-turn conversation reconstruction (production ReAct format)
#
# Every TenaOS task pipeline (form, report, cds, patient_education, and all
# scribe variants) is a genuine multi-turn ReAct loop in production: the
# runtime appends a live tool result as a new turn and makes a FRESH model
# call to decide what happens next (see form_pipeline/ciel_resolution.py,
# report_pipeline/ciel_resolution.py, tool_loop.py::KbAgentLoop,
# material_loop.py::PatientMaterialLoop, scribe_tool_loop.py::SoapScribeToolLoop).
#
# The raw traces captured under final_dataset/ already contain genuine
# per-turn event logs from real production/Vertex runs of those loops. This
# section reconstructs the actual multi-turn `conversations` array from that
# event log instead of flattening the whole run into one giant single-turn
# completion (the historical behavior of sft_example_from_trace below, which
# is why the shared LoRA adapter never saw a single real multi-turn example
# across its whole 16,005-example corpus).
#
# KNOWN LIMITATION: the event logs were built for UI display, not exact
# training reconstruction. Two gaps are known and are an explicit, documented
# trade-off rather than a silent approximation:
#   1. Pipeline-internal steering messages injected mid-loop (e.g. "You have
#      gathered sufficient KB evidence, call format_... now", the repeat-
#      search hint, budget nudges) are not separately journaled as their own
#      trace event in every pipeline, so they are not replayed as synthetic
#      user turns here. The tool-call -> live-tool-result -> next-decision
#      sequence (the actual behavior this retrain targets) IS fully captured.
#   2. tool_loop.py (CDS) and material_loop.py (patient education) do not
#      call `trace.add(...)` for every rejected/deduped tool_call turn (see
#      the empty-query / duplicate-query branches in KbAgentLoop.run), so a
#      strict minority of turns are invisible to this reconstruction. A more
#      complete fix would instrument those loops to log 100% of turns (or
#      log the literal `messages` array) rather than reconstructing post-hoc;
#      tracked as follow-up, out of scope for this pass.


@lru_cache(maxsize=None)
def _read_prompt_file(name: str) -> str:
    path = TENA_AGENT_PROMPTS_DIR / name
    if not path.exists():
        raise FileNotFoundError(
            f"Expected TenaAgent prompt file at {path}. Set TENA_AGENT_SERVICE_ROOT "
            "to the TenaAgent/service checkout if it lives elsewhere."
        )
    return path.read_text(encoding="utf-8").strip()


def _form_research_system() -> str:
    return _read_prompt_file("form_research_worklist_system.txt")


def _form_resolve_system() -> str:
    return _read_prompt_file("form_resolve_system.txt")


def _cds_system() -> str:
    return _read_prompt_file("cds_system.txt")


def _material_system() -> str:
    return _read_prompt_file("material_system.txt")


def _scribe_system() -> str:
    return _read_prompt_file("scribe_system.txt")


# Mirrors report_pipeline/planning_phase.py::_system_prompt() and
# report_pipeline/ciel_resolution.py::_resolve_system() exactly (both are
# inline Python string literals in that module, not prompt files, so they
# are copied here verbatim rather than read from disk). Keep in sync with
# those two functions if the report prompts change.
_REPORT_PLANNING_SYSTEM = (
    "You are a clinical report planner for OpenMRS. Convert the clinician request into a structured "
    "CIEL-backed report plan. Filters must be clinical concepts only. Sex, gender, age group, month, "
    "date, and time are groupBy dimensions, never filters. For coded requests like 'HIV positive', "
    "set valueKind='coded' and valueLabel='Positive'. For symptoms/findings, use valueKind='presence'. "
    "For numeric thresholds, set valueKind='numeric', operator, and numericThreshold. Indicators need "
    "a denominator; prefer encounters_in_range unless the user names a clinical denominator. Output by "
    "calling finalize_report_plan only. Treat common typos in date phrases, e.g. 'papst 12 months' "
    "means 'past 12 months'. Never create filters for words like past, last, month, line graph, trend, "
    "or chart. For 'month over month' or 'monthly' reports, use reportType='pivot', groupBy date_month, "
    "and time_series_line when the user asks for a line graph."
)

_REPORT_RESOLVE_SYSTEM = (
    "You are Gemma 4 running a CIEL-backed OpenMRS report-builder tool loop. "
    "Follow the typed plan exactly. First call get_report_draft. Set report type, date range, "
    "join mode, visualization, denominator/group_by. For each clinical filter, search CIEL, "
    "expand coded concepts when a value is needed, and add_filter with numeric CIEL ids only. "
    "For broad diagnosis requests like 'malaria diagnosis' or 'patients diagnosed with X', "
    "you MUST call search_related_ciel_concepts after the initial CIEL search to identify "
    "clinically related/narrower CIEL Diagnosis concepts. If the user's intent is broad, "
    "pass the relevant returned concepts together in "
    "add_filter as conceptIds so the filter is one logical OR over those CIEL concepts. "
    "Do not stop at the generic parent concept when search results include relevant subtype "
    "diagnoses with the same disease name, such as severe, uncomplicated, falciparum, cerebral, "
    "vivax, or complicated variants; include all clearly relevant non-retired Diagnosis-class "
    "CIEL concepts returned by the search. "
    "Coded filters need valueConceptId from the answers. Boolean/presence filters default true. "
    "Numeric filters need operator and numericThreshold. Build once, fix validation errors, then run once."
)


def _tena_call_text(name: str, arguments: dict[str, Any]) -> str:
    """Exact wire format production replays into assistant turn history.

    Matches form_pipeline/_llm_utils.py::assistant_message_for_tool_calls and
    tool_loop.py's text-tool fallback: a bare <tena_call>{...}</tena_call>
    wrapper, no surrounding prose (the corpus overwhelmingly used this form —
    e.g. toolCallId values of the form "text_tool_N" throughout the raw
    traces — rather than native OpenAI tool_calls).
    """
    return f'<tena_call>{json.dumps({"name": name, "arguments": arguments}, ensure_ascii=False)}</tena_call>'


def _assistant_tool_turn(name: str, arguments: dict[str, Any], *, think: str = "") -> dict[str, str]:
    content = _tena_call_text(name, arguments)
    if think:
        content = f"<think>{think}</think>\n{content}"
    return {"role": "assistant", "content": content}


def _tool_result_turn(name: str, tool_call_id: str, result: Any) -> dict[str, str]:
    """Matches llm_client.py::_strip_native_tool_messages exactly.

    EMPIRICALLY VERIFIED against the live production GGUF+chat_template (see
    the template-check step): every one of these ReAct loops calls
    llm.chat(..., tools=<schema>), and _messages_with_text_tool_instructions
    runs unconditionally whenever `tools` is not None — which converts EVERY
    role="tool" message into role="user" with content
    f"Result for {name}: {content}" BEFORE it ever reaches llama.cpp. A
    role="tool" message never actually reaches the wire in this system. An
    earlier version of this builder emitted role="tool" turns directly; a
    live llama-server --jinja dry run of that version showed the merged
    GGUF's chat template silently drops role="tool" content when the
    preceding assistant turn used the <tena_call> text wrapper (no native
    tool_calls array) rather than reconstructing this exact wire shape.
    """
    return {
        "role": "user",
        "content": f"Result for {name}: {json.dumps(result, ensure_ascii=False, default=str, sort_keys=True)}",
    }


def _looks_like_genuine_prose(text: str) -> bool:
    """Filter out leaked tool-call/JSON artifacts from agent_reasoning text.

    Some reasoning events capture a truncated <tena_call>...</tena_call> or
    raw JSON payload rather than real chain-of-thought prose (see
    research_phase.py's emit_thinking / extract_think_text — a closing tag
    that never arrives within the truncation window leaves the opening tag
    and raw JSON in the "text" field). Training on that leaked text as if it
    were legitimate reasoning would reinforce the same recital problem this
    retrain is meant to fix, so it is dropped rather than replayed.
    """
    stripped = text.strip()
    if not stripped:
        return False
    if stripped.startswith("{") or stripped.startswith("["):
        return False
    if "tena_call" in stripped or "tool_call" in stripped:
        return False
    return True


def _form_report_tool_turns(
    events: list[dict[str, Any]], *, phase: str | None = None
) -> list[dict[str, str]]:
    """Reconstruct assistant/tool turns from a form/report `events` list.

    Event shape: {"actor": "gemma"|"middleware"|"system", "operation": str,
    "payload": {...}}. A `model_tool_call` (payload.toolName/.arguments) is
    always immediately followed (possibly after a same-named middleware
    logging duplicate) by a `tool_result` (payload.result) — see
    ciel_resolution.py's store.append_event pairing.
    """
    turns: list[dict[str, str]] = []
    pending_think = ""
    i = 0
    while i < len(events):
        event = events[i]
        operation = str(event.get("operation") or "")
        payload = event.get("payload") if isinstance(event.get("payload"), dict) else {}
        if phase is not None and payload.get("phase") not in (phase, None):
            i += 1
            continue
        if event.get("actor") == "gemma" and operation == "agent_reasoning":
            text = str(payload.get("text") or "")
            if _looks_like_genuine_prose(text):
                pending_think = text
            i += 1
            continue
        if operation == "model_tool_call":
            name = str(payload.get("toolName") or "")
            arguments = payload.get("arguments") if isinstance(payload.get("arguments"), dict) else {}
            tool_call_id = str(payload.get("toolCallId") or f"call_{len(turns)}")
            # Find the paired tool_result, skipping same-named middleware
            # logging duplicates (e.g. "search_ciel_seeds", "update_form_draft").
            result = None
            j = i + 1
            while j < len(events):
                nxt = events[j]
                nxt_payload = nxt.get("payload") if isinstance(nxt.get("payload"), dict) else {}
                if nxt.get("operation") == "tool_result":
                    result = nxt_payload.get("result")
                    break
                if nxt.get("operation") in {"model_tool_call", "model_call"}:
                    break
                j += 1
            if name:
                turns.append(_assistant_tool_turn(name, arguments, think=pending_think))
                turns.append(_tool_result_turn(name, tool_call_id, result if result is not None else {}))
            pending_think = ""
            i = j + 1 if result is not None else i + 1
            continue
        i += 1
    return turns


def _final_text_turn(events: list[dict[str, Any]], *, phase: str, fallback: str) -> dict[str, str] | None:
    """Last genuine model prose turn for a form/report phase (the wrap-up
    summary after the final tool call resolves), if one was captured."""
    for event in reversed(events):
        payload = event.get("payload") if isinstance(event.get("payload"), dict) else {}
        if (
            event.get("actor") == "gemma"
            and event.get("operation") == "agent_reasoning"
            and payload.get("phase") == phase
        ):
            text = str(payload.get("text") or "")
            if _looks_like_genuine_prose(text):
                return {"role": "assistant", "content": text}
    if fallback:
        return {"role": "assistant", "content": fallback}
    return None


def _form_examples(record: dict[str, Any]) -> list[dict[str, Any]]:
    """Split into two standalone examples: research phase and resolve phase.

    Production runs these as two independent `llm.chat` conversations (each
    with its own system prompt and empty history) — run_research_phase and
    run_ciel_resolution never share a messages list — so training data must
    mirror that rather than merging both phases under one system prompt.
    """
    events = record.get("events") if isinstance(record.get("events"), list) else []
    source_request = record.get("source_request") if isinstance(record.get("source_request"), dict) else {}
    request = str(source_request.get("request") or "")
    trace_id = str(record.get("trace_id") or "")

    examples: list[dict[str, Any]] = []

    # --- Research phase: mirrors form_pipeline/research_phase.py::_research_user ---
    research_turns = _form_report_tool_turns(events, phase="research")
    if research_turns:
        research_user = (
            "mode: create\n"
            f"draft name: {trace_id}\n"
            f"form request: {request}\n\n"
            "current questions:\n(none yet)\n\n"
            "Research the clinical subject with WHO/MSF guideline searches, then call "
            "finalize_worklist with the questions to collect (do not repeat questions "
            "already present above)."
        )
        examples.append(
            _make_example(
                trace_id=f"{trace_id}_research",
                kind="form",
                schema_version=record.get("schema_version"),
                system=_form_research_system(),
                first_user=research_user,
                turns=research_turns,
                tool_schemas=_FORM_RESEARCH_TOOLS,
            )
        )

    # --- Resolve phase: mirrors form_pipeline/ciel_resolution.py::_resolve_user ---
    resolve_turns = _form_report_tool_turns(events, phase="resolve")
    if resolve_turns:
        final_turn = _final_text_turn(events, phase="resolve", fallback="")
        if final_turn:
            resolve_turns = resolve_turns + [final_turn]
        resolve_user = (
            f"draftId: {trace_id}\n"
            f"user request: {request}\n\n"
            "=== plan (resolve each question in order) ===\n(see prior research phase)\n\n"
            "=== current basket ===\nNo questions yet.\n\n"
            "Resolve each planned question against CIEL, then send ONE update_form_draft "
            "with all add_section + add_field operations and call build_form_schema. "
            "When the basket reflects the plan, return a short final summary."
        )
        examples.append(
            _make_example(
                trace_id=f"{trace_id}_resolve",
                kind="form",
                schema_version=record.get("schema_version"),
                system=_form_resolve_system(),
                first_user=resolve_user,
                turns=resolve_turns,
                tool_schemas=_FORM_RESOLVE_TOOLS,
            )
        )
    return examples


def _report_examples(record: dict[str, Any]) -> list[dict[str, Any]]:
    """Split into planning phase and resolve phase (see _form_examples)."""
    events = record.get("events") if isinstance(record.get("events"), list) else []
    source_request = record.get("source_request") if isinstance(record.get("source_request"), dict) else {}
    request = str(source_request.get("request") or "")
    trace_id = str(record.get("trace_id") or "")

    examples: list[dict[str, Any]] = []

    # The planning phase's finalize_report_plan call is never journaled as a
    # model_tool_call — planning_phase.py::_emit_plan_event logs the already-
    # parsed plan dict directly on a "report_plan"-phase agent_reasoning event
    # instead of replaying the tool call. Reconstruct the equivalent turn
    # from that payload.
    planning_turns: list[dict[str, str]] = []
    for event in events:
        payload = event.get("payload") if isinstance(event.get("payload"), dict) else {}
        if event.get("actor") == "gemma" and payload.get("phase") == "report_plan" and isinstance(payload.get("plan"), dict):
            planning_turns.append(_assistant_tool_turn("finalize_report_plan", payload["plan"]))
            break
    if planning_turns:
        planning_user = (
            "mode: create\n"
            f"report name: {trace_id}\n"
            f"user request: {request}\n\n"
            "current spec:\n{}\n\n"
            "Plan only the requested report/change. Do not invent CIEL ids."
        )
        examples.append(
            _make_example(
                trace_id=f"{trace_id}_planning",
                kind="report",
                schema_version=record.get("schema_version"),
                system=_REPORT_PLANNING_SYSTEM,
                first_user=planning_user,
                turns=planning_turns,
                tool_schemas=_REPORT_PLAN_TOOLS,
            )
        )

    resolve_turns = _form_report_tool_turns(events, phase="report_resolve")
    if resolve_turns:
        final_turn = _final_text_turn(events, phase="report_resolve", fallback="")
        if final_turn:
            resolve_turns = resolve_turns + [final_turn]
        resolve_user = (
            f"draftId: {trace_id}\n"
            f"user request: {request}\n\n"
            "=== typed report plan ===\n(see prior planning phase)\n\n"
            "=== current spec ===\n{}\n\n"
            "Use only report tools. Do not invent CIEL ids or result numbers."
        )
        examples.append(
            _make_example(
                trace_id=f"{trace_id}_resolve",
                kind="report",
                schema_version=record.get("schema_version"),
                system=_REPORT_RESOLVE_SYSTEM,
                first_user=resolve_user,
                turns=resolve_turns,
                tool_schemas=_REPORT_RESOLVE_TOOLS,
            )
        )
    return examples


#: trace_id -> ordered list of reconstructed search results, one per genuine
#: (non-empty, non-duplicate) search_guidelines call in that trace, each
#: shaped exactly like tool_loop.py/material_loop.py's real wire response:
#: {"hits": [_compact_hit(h), ...], "total_accumulated": N}. Populated by
#: reconstruct_kb_hits.py, which re-runs each trace's recorded queries
#: against the live who_msf_guidelines KB (the traces themselves never
#: journaled the full wire payload -- only a UI-summary -- so this is a
#: reconstruction via live re-query, not a byte-identical replay; see that
#: script's docstring for exact parity details and caveats).
_KB_HITS_CACHE_PATH = ROOT / "lora_training" / "artifacts_v2" / "kb_hits_reconstructed.json"
_KB_HITS_CACHE: dict[str, list[dict[str, Any]]] | None = None
_KB_HITS_CACHE_MISSES = 0


def _load_kb_hits_cache() -> dict[str, list[dict[str, Any]]]:
    global _KB_HITS_CACHE
    if _KB_HITS_CACHE is None:
        if _KB_HITS_CACHE_PATH.exists():
            _KB_HITS_CACHE = json.loads(_KB_HITS_CACHE_PATH.read_text(encoding="utf-8"))
        else:
            _KB_HITS_CACHE = {}
    return _KB_HITS_CACHE


def _reconstructed_search_result(trace_id: str, search_ordinal: int) -> dict[str, Any] | None:
    """Look up the reconstructed full wire-format hits for the `search_ordinal`-th
    (0-indexed) genuine search_guidelines call in `trace_id`. Returns None on
    any cache miss so the caller can fall back to the UI-summary payload.
    """
    global _KB_HITS_CACHE_MISSES
    cache = _load_kb_hits_cache()
    entries = cache.get(trace_id)
    if entries is None or search_ordinal >= len(entries):
        _KB_HITS_CACHE_MISSES += 1
        return None
    entry = entries[search_ordinal]
    if not isinstance(entry, dict) or "hits" not in entry:
        _KB_HITS_CACHE_MISSES += 1
        return None
    return entry


def _kb_agent_loop_turns(
    events: list[dict[str, Any]],
    *,
    trace_id: str,
    min_searches: int,
    steering_message: Callable[[int, list[str]], str],
) -> tuple[list[dict[str, str]], dict[str, Any] | None]:
    """Reconstruct assistant/tool turns from a cds/patient_education
    `trace.events` list (shape: {"type", "title", "detail", "payload"}).

    Returns (turns, final_format_call_arguments_or_None). Tool calls carry
    their arguments directly in payload["arguments"] (native tool_calls were
    used for search_guidelines in these two loops — see tool_loop.py /
    material_loop.py).

    Three behaviors from tool_loop.py::KbAgentLoop.run / material_loop.py's
    equivalent that were previously dropped are reconstructed here (see the
    "Training vs Production ReAct Loop" audit, 2026-07-02):

      1. search_guidelines results now use the full {"hits": [...],
         "total_accumulated": N} wire payload (via _reconstructed_search_result)
         instead of the UI-summary {"query", "hits_returned", "top_hit"} the
         `middleware_result` event actually journals. Falls back to the
         UI-summary shape (with a QA-visible miss counter) if the
         reconstruction cache has no entry for a given call.
      2. `model_reasoning` events capture the "think out loud" prose the CDS/
         patient-education system prompts require before each search call;
         the immediately-preceding, not-yet-consumed one is attached to the
         following assistant tool turn as its <think> block.
      3. Whenever a completed, counted search still leaves search_count below
         the production minimum, tool_loop.py/material_loop.py append a
         synthetic steering "role": "user" turn nudging the next search --
         reconstructed deterministically here from search_count/
         searched_queries, which are already known at this point. This only
         covers the AFTER-A-COMPLETED-SEARCH steering branch; the separate
         "model returned no tool call at all" branch is not reconstructable
         because such turns leave no `model_tool_call` event in the trace at
         all (see the module-level KNOWN LIMITATION note above).
    """
    turns: list[dict[str, str]] = []
    final_args: dict[str, Any] | None = None
    search_count = 0
    searched_queries: list[str] = []
    pending_think = ""
    i = 0
    call_index = 0
    while i < len(events):
        event = events[i]
        etype = str(event.get("type") or "")
        payload = event.get("payload") if isinstance(event.get("payload"), dict) else {}
        if etype == "model_reasoning":
            text = str(event.get("detail") or "")
            if text:
                pending_think = text
            i += 1
            continue
        if etype == "model_tool_call":
            name = str(event.get("title") or "")
            arguments = payload.get("arguments") if isinstance(payload.get("arguments"), dict) else {}
            tool_call_id = f"call_{call_index}"
            call_index += 1
            think = pending_think
            pending_think = ""
            if name in {"format_cds_result", "format_patient_material"}:
                final_args = arguments
                i += 1
                continue
            result: Any = {}
            j = i + 1
            while j < len(events):
                nxt = events[j]
                if str(nxt.get("type") or "") == "middleware_result":
                    nxt_payload = nxt.get("payload") if isinstance(nxt.get("payload"), dict) else {}
                    result = nxt_payload if nxt_payload else str(nxt.get("detail") or "")
                    break
                if str(nxt.get("type") or "") == "model_tool_call":
                    break
                j += 1
            is_new_search = False
            if name == "search_guidelines":
                norm_query = " ".join(str(arguments.get("query") or "").lower().split())
                is_new_search = bool(norm_query) and norm_query not in searched_queries
                if is_new_search:
                    reconstructed = _reconstructed_search_result(trace_id, search_count)
                    if reconstructed is not None:
                        result = reconstructed
                    searched_queries.append(norm_query)
                    search_count += 1
            turns.append(_assistant_tool_turn(name, arguments, think=think))
            turns.append(_tool_result_turn(name, tool_call_id, result))
            if is_new_search and search_count < min_searches:
                turns.append({"role": "user", "content": steering_message(search_count, searched_queries)})
            i = j + 1
            continue
        i += 1
    return turns, final_args


def _cds_steering_message(search_count: int, searched_queries: list[str]) -> str:
    """Mirrors tool_loop.py::KbAgentLoop.run's post-search steering nudge
    (the branch fired after a completed, counted search_guidelines call
    while still below _MIN_SEARCHES) verbatim."""
    return (
        f"You have completed {search_count} unique KB search(es). "
        f"Previous queries: {', '.join(sorted(searched_queries)) or 'none'}. "
        "Review the latest result, identify the next clinical gap, then call "
        "search_guidelines exactly once with a NEW focused query."
    )


def _cds_examples(record: dict[str, Any]) -> list[dict[str, Any]]:
    trace = record.get("trace") if isinstance(record.get("trace"), dict) else {}
    events = trace.get("events") if isinstance(trace.get("events"), list) else []
    trace_id = str(record.get("trace_id") or "")
    turns, final_args = _kb_agent_loop_turns(
        events,
        trace_id=trace_id,
        min_searches=_CDS_MIN_SEARCHES,
        steering_message=_cds_steering_message,
    )
    if not turns:
        return []
    context_summary = ""
    for event in events:
        if str(event.get("type") or "") == "context":
            context_summary = str(((event.get("payload") or {}).get("summary")) or "")
            break
    structured = record.get("structured_cds") if isinstance(record.get("structured_cds"), dict) else {}
    had_real_format_call = final_args is not None
    if final_args is None:
        final_args = {
            "status": structured.get("status") or "recommendation",
            "summary": structured.get("summary") or "",
            "content": structured.get("content") or "",
        }
    if had_real_format_call:
        # tool_loop.py::KbAgentLoop.run always injects this exact nudge as the
        # last user turn immediately before the model's format_cds_result call
        # (see "Call format_cds_result now" branch, fired once search_count >=
        # _MIN_SEARCHES). Only reconstruct it when we know a real format call
        # was actually captured in the trace -- if final_args was synthesized
        # from structured_cds (loop_exhausted fallback with no model call),
        # production never sent this message either.
        turns.append({
            "role": "user",
            "content": (
                "You have gathered sufficient KB evidence. "
                "Call format_cds_result now to emit the structured CDS report."
            ),
        })
    turns = turns + [_assistant_tool_turn("format_cds_result", final_args)]
    user = (
        f"Patient context:\n{context_summary}\n\n"
        "Begin Phase A: read the patient context carefully. Identify the primary diagnosis, "
        "ALL active medications, key comorbidities, and risk factors. Plan your search queries. "
        "Then issue your FIRST search_guidelines call with a focused 3-6 word query "
        "(e.g. 'hypertension first-line treatment', 'diabetes management guidelines')."
    )
    return [
        _make_example(
            trace_id=trace_id,
            kind="cds",
            schema_version=record.get("schema_version"),
            system=_cds_system(),
            first_user=user,
            turns=turns,
            tool_schemas=_CDS_TOOL_SCHEMAS,
        )
    ]


def _patient_education_steering_message(search_count: int, searched_queries: list[str]) -> str:
    """Mirrors material_loop.py's post-search steering nudge (the branch
    fired after a completed, counted search_guidelines call while still
    below its inline `search_count < 4` threshold) verbatim."""
    return (
        f"You have completed {search_count} unique KB search(es). "
        f"Previous queries: {', '.join(sorted(searched_queries)) or 'none'}. "
        "You need at least 4 unique searches before writing the full material. "
        "Review the latest result, identify the single next gap, then call search_guidelines exactly once with a NEW focused query. "
        "Do not batch multiple searches."
    )


def _patient_education_examples(record: dict[str, Any]) -> list[dict[str, Any]]:
    trace = record.get("trace") if isinstance(record.get("trace"), dict) else {}
    events = trace.get("events") if isinstance(trace.get("events"), list) else []
    trace_id = str(record.get("trace_id") or "")
    turns, final_args = _kb_agent_loop_turns(
        events,
        trace_id=trace_id,
        min_searches=_MATERIAL_MIN_SEARCHES,
        steering_message=_patient_education_steering_message,
    )
    if not turns:
        return []
    context_summary = ""
    for event in events:
        if str(event.get("type") or "") == "context":
            context_summary = str(((event.get("payload") or {}).get("summary")) or "")
            break
    material = record.get("material") if isinstance(record.get("material"), dict) else {}
    had_real_format_call = final_args is not None
    if final_args is None:
        final_args = {"title": material.get("title") or "", "content": material.get("content") or ""}
    if had_real_format_call:
        # material_loop.py's equivalent forced-format nudge, injected
        # immediately before the model's format_patient_material call once
        # search_count >= min_searches. Same caveat as CDS: only reconstruct
        # when a real model format call was actually captured in the trace.
        turns.append({
            "role": "user",
            "content": (
                "You have gathered sufficient KB evidence. "
                "Now call format_patient_material to emit the complete 7-section patient education material."
            ),
        })
    turns = turns + [_assistant_tool_turn("format_patient_material", final_args)]
    user = (
        f"Patient information:\n{context_summary}\n\n"
        "Begin Phase A: read the patient information above, identify the primary diagnosis, "
        "all active medications, key risk factors, and the likely information gaps. "
        "Then issue exactly ONE search_guidelines call: your FIRST short, specific query "
        "(3-7 words, e.g. 'hypertension blood pressure causes'). Do not call the other searches yet."
    )
    return [
        _make_example(
            trace_id=trace_id,
            kind="patient_education",
            schema_version=record.get("schema_version"),
            system=_material_system(),
            first_user=user,
            turns=turns,
            tool_schemas=_MATERIAL_TOOL_SCHEMAS,
        )
    ]


_TENA_CALL_RE = re.compile(r"<(?:tena_call|tool_call)>\s*(.*?)\s*</(?:tena_call|tool_call)>", re.DOTALL)


def _parse_scribe_tool_call(detail: str) -> tuple[str, dict[str, Any]] | None:
    match = _TENA_CALL_RE.search(detail or "")
    if not match:
        return None
    try:
        obj = json.loads(match.group(1))
    except json.JSONDecodeError:
        return None
    if not isinstance(obj, dict):
        return None
    name = obj.get("name")
    if not isinstance(name, str) or not name:
        return None
    arguments = obj.get("arguments") if isinstance(obj.get("arguments"), dict) else {}
    return name, arguments


def _scribe_loop_turns(generation_trace: list[dict[str, Any]]) -> list[dict[str, str]]:
    """Reconstruct assistant/tool turns from SoapScribeToolLoop's
    `result.generationTrace` (shape: {"type", "title", "detail", "payload"}),
    where the tool call itself is embedded as <tena_call>{...}</tena_call>
    text inside a `model_reasoning` event's `detail` field rather than a
    separate structured field — see scribe_tool_loop.py::_add_trace.

    Some traces (observed for voice_scribe_audio) skip the ReAct tool loop
    entirely and have the model return one bare JSON object with a "soap"
    key directly as its whole turn (parse_scribe_response's legacy-JSON
    path) — no <tena_call> wrapper at all. That is treated as a direct
    finalize_soap_note turn with those exact arguments rather than forced
    into the tool-call shape it never actually used.
    """
    turns: list[dict[str, str]] = []
    call_index = 0
    i = 0
    while i < len(generation_trace):
        event = generation_trace[i]
        etype = str(event.get("type") or "")
        if etype == "model_reasoning":
            detail = str(event.get("detail") or "")
            parsed = _parse_scribe_tool_call(detail)
            if parsed is not None:
                name, arguments = parsed
                tool_call_id = f"call_{call_index}"
                call_index += 1
                result: Any = {}
                j = i + 1
                while j < len(generation_trace):
                    nxt = generation_trace[j]
                    nxt_type = str(nxt.get("type") or "")
                    if nxt_type == "middleware_result":
                        nxt_payload = nxt.get("payload") if isinstance(nxt.get("payload"), dict) else {}
                        result = nxt_payload if nxt_payload else str(nxt.get("detail") or "")
                        break
                    if nxt_type == "model_reasoning" and _parse_scribe_tool_call(str(nxt.get("detail") or "")):
                        break
                    j += 1
                turns.append(_assistant_tool_turn(name, arguments))
                turns.append(_tool_result_turn(name, tool_call_id, result))
                i = j + 1
                continue
            try:
                bare = json.loads(detail)
            except json.JSONDecodeError:
                bare = None
            if isinstance(bare, dict) and "soap" in bare:
                turns.append(_assistant_tool_turn("finalize_soap_note", bare))
                i += 1
                continue
        i += 1
    return turns


def _scribe_pass_example(
    *, trace_id: str, note_text: str, patient_summary: str, result: dict[str, Any], schema_version: Any
) -> dict[str, Any] | None:
    generation_trace = result.get("generationTrace") if isinstance(result.get("generationTrace"), list) else []
    turns = _scribe_loop_turns(generation_trace)
    if not turns:
        return None
    # Final finalize_soap_note call, if not already the last reconstructed turn
    # (a rejected premature finalize earlier in the trace is a real retry turn,
    # already captured above; this appends the one that actually stuck).
    final_args = {
        "soap": result.get("soap") or {},
        "concepts": result.get("concepts") or [],
        "observations": result.get("observations") or [],
        "medications": result.get("medications") or [],
    }
    last_assistant_content = next((t["content"] for t in reversed(turns) if t["role"] == "assistant"), "")
    last_is_finalize = "finalize_soap_note" in last_assistant_content
    if not last_is_finalize:
        turns = turns + [_assistant_tool_turn("finalize_soap_note", final_args)]
    content = ""
    if patient_summary:
        content += f"[BACKGROUND - DO NOT EXTRACT FROM THIS]:\n{patient_summary}\n\n"
    content += (
        f"[NOTE TO SCRIBE - EXTRACT ONLY FROM THIS]:\n\"\"\"\n{note_text.strip()}\n\"\"\"\n\n"
        "Begin by searching CIEL for the Assessment diagnosis or Plan medication. "
        "Then finalize SOAP and saveable items."
    )
    return _make_example(
        trace_id=trace_id,
        kind="scribe",
        schema_version=schema_version,
        system=_scribe_system(),
        first_user=content,
        turns=turns,
    )


def _scribe_text_english_examples(record: dict[str, Any]) -> list[dict[str, Any]]:
    sample = record.get("source_sample") if isinstance(record.get("source_sample"), dict) else {}
    result = record.get("result") if isinstance(record.get("result"), dict) else {}
    example = _scribe_pass_example(
        trace_id=str(record.get("trace_id") or ""),
        note_text=str(sample.get("noteText") or ""),
        patient_summary=str(sample.get("patientSummary") or ""),
        result=result,
        schema_version=record.get("schema_version"),
    )
    return [example] if example else []


def _scribe_text_amharic_examples(record: dict[str, Any]) -> list[dict[str, Any]]:
    source_pair = record.get("source_pair") if isinstance(record.get("source_pair"), dict) else {}
    english_trace = record.get("english_trace") if isinstance(record.get("english_trace"), dict) else {}
    trace_id = str(record.get("trace_id") or "")
    examples: list[dict[str, Any]] = []
    # Amharic pass: the actual target behavior for the [scribe-am] tag.
    example = _scribe_pass_example(
        trace_id=f"{trace_id}_am",
        note_text=str(source_pair.get("noteText") or ""),
        patient_summary="",
        result=english_trace.get("result") if isinstance(english_trace.get("result"), dict) else {},
        schema_version=record.get("schema_version"),
    )
    if example:
        examples.append(example)
    return examples


def _voice_scribe_audio_examples(record: dict[str, Any]) -> list[dict[str, Any]]:
    source_audio = record.get("source_audio") if isinstance(record.get("source_audio"), dict) else {}
    audio_trace = record.get("audio_trace") if isinstance(record.get("audio_trace"), dict) else {}
    result = audio_trace.get("result") if isinstance(audio_trace.get("result"), dict) else {}
    # NOTE: the true ASR transcript fed to the loop at generation time isn't
    # preserved in the canonical trace; source_audio.noteText (the reference
    # note text used to synthesize the audio) is the closest available proxy
    # and is used here as a documented approximation.
    example = _scribe_pass_example(
        trace_id=str(record.get("trace_id") or ""),
        note_text=str(source_audio.get("noteText") or ""),
        patient_summary=str(source_audio.get("patientSummary") or ""),
        result=result,
        schema_version=record.get("schema_version"),
    )
    return [example] if example else []


def _make_example(
    *,
    trace_id: str,
    kind: str,
    schema_version: Any,
    system: str,
    first_user: str,
    turns: list[dict[str, str]],
    tool_schemas: list[dict[str, Any]] | None = None,
) -> dict[str, Any]:
    conversations: list[dict[str, str]] = []
    if tool_schemas is not None:
        # Mirrors production wire order exactly: [tool-schema system][task
        # system][user]... -- see the module-level note near
        # _tool_schema_system_message. Pass tool_schemas=None only for kinds
        # that never call llm.chat(..., tools=...) (the scribe variants).
        conversations.append(_tool_schema_message_for(tool_schemas))
    conversations.append({"role": "system", "content": system})
    conversations.append({"role": "user", "content": first_user})
    conversations.extend(turns)
    return {
        "id": trace_id,
        "kind": kind,
        "task_tag": TASK_TAGS.get(kind, ""),
        "schema_version": schema_version,
        "conversations": conversations,
    }


_MULTITURN_BUILDERS: dict[str, Callable[[dict[str, Any]], list[dict[str, Any]]]] = {
    "form": _form_examples,
    "report": _report_examples,
    "cds": _cds_examples,
    "patient_education": _patient_education_examples,
    "scribe_text_english": _scribe_text_english_examples,
    "scribe_text_amharic": _scribe_text_amharic_examples,
    "voice_scribe_audio": _voice_scribe_audio_examples,
}


def multiturn_examples_from_trace(kind: str, record: dict[str, Any]) -> list[dict[str, Any]]:
    """Build one or more real multi-turn SFT examples from a trace's event log.

    Returns a list because two-phase pipelines (form: research+resolve;
    report: planning+resolve) are split into two independent examples, one
    per phase, mirroring the two independent `llm.chat` conversations
    production actually runs (see _form_examples/_report_examples docstrings).
    Returns an empty list if the trace's event log could not be turned into
    any usable multi-turn conversation (caller falls back to the legacy
    single-completion format for that record only, and this is counted and
    surfaced rather than silently swallowed — see build_sft's fallback_counts).
    """
    builder = _MULTITURN_BUILDERS.get(kind)
    if builder is None:
        return []
    return builder(record)


def build_sft(out_dir: Path, *, seed: int = 3407) -> dict[str, Any]:
    normalized_root = out_dir / "normalized"
    sft_root = out_dir / "sft"
    rows_by_split: dict[str, list[dict[str, Any]]] = {"train": [], "validation": [], "test": []}
    per_kind_counts: dict[str, Counter[str]] = {}
    fallback_counts: dict[str, int] = {}
    for kind in sorted(TRACE_KINDS):
        path = normalized_root / kind / f"{kind}_accepted.jsonl"
        if not path.exists():
            continue
        split_counts: Counter[str] = Counter()
        for _, record, error in iter_jsonl(path):
            if error or record is None:
                continue
            split = split_for_record(record, seed=seed)
            examples = multiturn_examples_from_trace(kind, record)
            if not examples:
                # Explicit, counted fallback — NOT a silent default. See the
                # module-level note above `multiturn_examples_from_trace`.
                fallback_counts[kind] = fallback_counts.get(kind, 0) + 1
                examples = [sft_example_from_trace(kind, record)]
            for example in examples:
                rows_by_split[split].append(example)
                split_counts.update([split])
        per_kind_counts[kind] = split_counts
    for split, rows in rows_by_split.items():
        rows.sort(key=lambda row: str(row["id"]))
        write_jsonl(sft_root / f"{split}.jsonl", rows)
    summary = {
        "schema_version": "tenaos_lora_sft_summary_v1",
        "created_at": utc_now(),
        "source_normalized_root": str(normalized_root),
        "output_root": str(sft_root),
        "seed": seed,
        "split_counts": {split: len(rows) for split, rows in rows_by_split.items()},
        "per_kind_split_counts": {kind: dict(counter) for kind, counter in sorted(per_kind_counts.items())},
        "single_completion_fallback_counts": fallback_counts,
        "format": (
            "Real multi-turn ShareGPT-style conversations reconstructed from production "
            "event traces (system/user/assistant tool-calls/tool results), one or two per "
            "source trace depending on whether the pipeline is single- or two-phase; falls "
            "back to a single-completion recital only for records where no usable event log "
            "was found (see single_completion_fallback_counts). Train script applies the "
            "Gemma 4 chat template and masks loss to assistant turns only."
        ),
    }
    write_json(sft_root / "sft_summary.json", summary)
    return summary


def split_for_record(record: dict[str, Any], *, seed: int) -> str:
    trace_id = str(record.get("trace_id") or "")
    digest = hashlib.sha256(f"{seed}:{trace_id}".encode("utf-8")).hexdigest()
    bucket = int(digest[:8], 16) % 100
    if bucket < 5:
        return "test"
    if bucket < 10:
        return "validation"
    return "train"


def sft_example_from_trace(kind: str, record: dict[str, Any]) -> dict[str, Any]:
    tag = TASK_TAGS[kind]
    user_payload = build_user_payload(kind, record)
    assistant_payload = build_assistant_payload(kind, record)
    return {
        "id": record.get("trace_id"),
        "kind": kind,
        "task_tag": tag,
        "schema_version": record.get("schema_version"),
        "conversations": [
            {
                "role": "user",
                "content": f"{tag}\nYou are TenaOS Gemma 4 E4B. Produce the validated clinical-informatics trace for this request.\n\n{user_payload}",
            },
            {
                "role": "assistant",
                "content": json.dumps(assistant_payload, ensure_ascii=False, sort_keys=True),
            },
        ],
    }


def build_user_payload(kind: str, record: dict[str, Any]) -> str:
    if kind == "form":
        source = record.get("source_request") if isinstance(record.get("source_request"), dict) else {}
        return json.dumps({"request": source.get("request"), "metadata": source.get("metadata")}, ensure_ascii=False, sort_keys=True)
    if kind == "report":
        source = record.get("source_request") if isinstance(record.get("source_request"), dict) else {}
        return json.dumps({"request": source.get("request"), "metadata": source.get("metadata")}, ensure_ascii=False, sort_keys=True)
    if kind == "scribe_text_english":
        source = record.get("source_sample") if isinstance(record.get("source_sample"), dict) else {}
        return json.dumps({"noteText": source.get("noteText"), "patientSummary": source.get("patientSummary"), "expected": source.get("expected")}, ensure_ascii=False, sort_keys=True)
    if kind == "voice_scribe_audio":
        source = record.get("source_audio") if isinstance(record.get("source_audio"), dict) else {}
        return json.dumps({"audio_path": source.get("audio_path"), "noteText_reference": source.get("noteText"), "expected": source.get("expected")}, ensure_ascii=False, sort_keys=True)
    if kind == "scribe_text_amharic":
        source = record.get("source_pair") if isinstance(record.get("source_pair"), dict) else {}
        return json.dumps({"noteText": source.get("noteText"), "source_example_id": source.get("source_example_id"), "expected": source.get("expected")}, ensure_ascii=False, sort_keys=True)
    if kind in {"cds", "patient_education"}:
        source = record.get("source_case") if isinstance(record.get("source_case"), dict) else {}
        return json.dumps({"case_id": source.get("case_id"), "metadata": source.get("metadata"), "patient_summary": source.get("patient_summary") or source.get("summary")}, ensure_ascii=False, sort_keys=True)
    return json.dumps({"trace_id": record.get("trace_id")}, ensure_ascii=False, sort_keys=True)


def build_assistant_payload(kind: str, record: dict[str, Any]) -> dict[str, Any]:
    # Keep source inputs out of the assistant target when possible; the target is
    # the observable, validated assistant-side behavior and final artifact.
    excluded = {
        "source_request",
        "source_sample",
        "source_audio",
        "source_pair",
        "source_case",
        "request_metadata",
        "sample_metadata",
        "audio_metadata",
        "case_metadata",
        "pair_metadata",
    }
    return {key: value for key, value in record.items() if key not in excluded}


def evaluate(out_dir: Path) -> dict[str, Any]:
    normalized_root = out_dir / "normalized"
    sft_root = out_dir / "sft"
    kind_results = []
    failures: list[str] = []
    for kind in sorted(TRACE_KINDS):
        path = normalized_root / kind / f"{kind}_accepted.jsonl"
        if not path.exists():
            failures.append(f"{kind}:missing_normalized_accepted")
            continue
        count = 0
        schema_bad = 0
        tag_bad = 0
        ready_bad = 0
        for _, record, error in iter_jsonl(path):
            if error or record is None:
                failures.append(f"{kind}:invalid_json")
                continue
            count += 1
            if record.get("schema_version") != CANONICAL_SCHEMA[kind]:
                schema_bad += 1
            if record.get("task_tag") != TASK_TAGS[kind]:
                tag_bad += 1
            quality = record.get("quality") if isinstance(record.get("quality"), dict) else {}
            if quality.get("ready_for_training") is not True:
                ready_bad += 1
        if schema_bad:
            failures.append(f"{kind}:schema_bad:{schema_bad}")
        if tag_bad:
            failures.append(f"{kind}:tag_bad:{tag_bad}")
        if ready_bad:
            failures.append(f"{kind}:ready_bad:{ready_bad}")
        kind_results.append(
            {
                "kind": kind,
                "accepted_count": count,
                "schema_bad": schema_bad,
                "tag_bad": tag_bad,
                "ready_bad": ready_bad,
            }
        )
    split_counts = {}
    for split in ("train", "validation", "test"):
        path = sft_root / f"{split}.jsonl"
        if not path.exists():
            failures.append(f"sft:{split}:missing")
            continue
        split_counts[split] = sum(1 for _, record, error in iter_jsonl(path) if record is not None and not error)
    result = {
        "schema_version": "tenaos_lora_corpus_eval_v1",
        "created_at": utc_now(),
        "passed": not failures,
        "failures": failures,
        "kind_results": kind_results,
        "split_counts": split_counts,
        "report_metric_placeholders": {
            "form_eval": "Run production held-out form recall/schema-valid/bad-code eval after adapter training.",
            "report_eval": "Run report plan correctness/count accuracy eval after adapter training.",
            "scribe_eval": "Run SOAP completeness and extraction F1 eval after adapter training.",
            "cds_education_eval": "Run citation, unsupported recommendation, abstention, and no-invented-dose eval after adapter training.",
        },
    }
    write_json(out_dir / "eval" / "corpus_eval_summary.json", result)
    return result


def package(out_dir: Path, *, adapter_path: Path | None = None, base_model: str = "unsloth/gemma-4-E4B-it") -> dict[str, Any]:
    normalized_summary_path = out_dir / "normalized" / "normalization_summary.json"
    sft_summary_path = out_dir / "sft" / "sft_summary.json"
    eval_summary_path = out_dir / "eval" / "corpus_eval_summary.json"
    metadata = {
        "schema_version": "tenaos_lora_adapter_metadata_v1",
        "created_at": utc_now(),
        "technical_report": str(REPORT_PATH),
        "base_model": base_model,
        "adapter_path": str(adapter_path) if adapter_path else "",
        "adapter_present": bool(adapter_path and adapter_path.exists()),
        "training_hyperparameters": {
            "lora_rank": 16,
            "lora_alpha": 32,
            "lora_dropout": 0.05,
            "target_modules": "attention_and_mlp_projections",
            "epochs": 3,
            "precision": "bf16_16bit_lora_default; optional 4bit fallback",
            "seed": 3407,
        },
        "dataset_artifacts": {
            "normalization_summary": str(normalized_summary_path),
            "sft_summary": str(sft_summary_path),
            "corpus_eval_summary": str(eval_summary_path),
            "train_jsonl": str(out_dir / "sft" / "train.jsonl"),
            "validation_jsonl": str(out_dir / "sft" / "validation.jsonl"),
            "test_jsonl": str(out_dir / "sft" / "test.jsonl"),
        },
        "artifact_hashes": artifact_hashes(out_dir / "sft"),
    }
    for key, path in (
        ("normalization", normalized_summary_path),
        ("sft", sft_summary_path),
        ("eval", eval_summary_path),
    ):
        metadata[f"{key}_summary"] = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {}
    package_dir = out_dir / "package"
    write_json(package_dir / "adapter_metadata.json", metadata)
    instructions = f"""# TenaOS LoRA Adapter Package

Base model: `{base_model}`

Train with:

```bash
python3 lora_training/train_unsloth_lora.py \\
  --train-jsonl {out_dir / "sft" / "train.jsonl"} \\
  --validation-jsonl {out_dir / "sft" / "validation.jsonl"} \\
  --output-dir {out_dir / "adapter"} \\
  --base-model {base_model}
```

The trainer saves the LoRA adapter, trainer metrics, and a metadata copy. Use
`--merge-output-dir` on the trainer command to export merged weights after a
successful training run.
"""
    (package_dir / "MERGE_INSTRUCTIONS.md").write_text(instructions, encoding="utf-8")
    return metadata


def artifact_hashes(path: Path) -> dict[str, str]:
    hashes = {}
    if not path.exists():
        return hashes
    for child in sorted(path.glob("*.jsonl")):
        digest = hashlib.sha256()
        with child.open("rb") as handle:
            for chunk in iter(lambda: handle.read(1024 * 1024), b""):
                digest.update(chunk)
        hashes[str(child)] = digest.hexdigest()
    return hashes


def run_all(final_dataset: Path, out_dir: Path, *, seed: int, adapter_path: Path | None, base_model: str) -> None:
    audit(final_dataset, out_dir)
    normalize(final_dataset, out_dir)
    build_sft(out_dir, seed=seed)
    evaluate(out_dir)
    package(out_dir, adapter_path=adapter_path, base_model=base_model)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest="command", required=True)

    def add_common(command: argparse.ArgumentParser) -> None:
        command.add_argument("--final-dataset", type=Path, default=DEFAULT_FINAL_DATASET)
        command.add_argument("--out-dir", type=Path, default=DEFAULT_WORK_DIR)

    for name in ("audit", "normalize", "evaluate"):
        add_common(sub.add_parser(name))

    build = sub.add_parser("build-sft")
    add_common(build)
    build.add_argument("--seed", type=int, default=3407)

    pkg = sub.add_parser("package")
    add_common(pkg)
    pkg.add_argument("--adapter-path", type=Path)
    pkg.add_argument("--base-model", default="unsloth/gemma-4-E4B-it")

    all_cmd = sub.add_parser("all")
    add_common(all_cmd)
    all_cmd.add_argument("--seed", type=int, default=3407)
    all_cmd.add_argument("--adapter-path", type=Path)
    all_cmd.add_argument("--base-model", default="unsloth/gemma-4-E4B-it")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.command == "audit":
        result = audit(args.final_dataset, args.out_dir)
    elif args.command == "normalize":
        result = normalize(args.final_dataset, args.out_dir)
    elif args.command == "build-sft":
        result = build_sft(args.out_dir, seed=args.seed)
    elif args.command == "evaluate":
        result = evaluate(args.out_dir)
    elif args.command == "package":
        result = package(args.out_dir, adapter_path=args.adapter_path, base_model=args.base_model)
    elif args.command == "all":
        run_all(args.final_dataset, args.out_dir, seed=args.seed, adapter_path=args.adapter_path, base_model=args.base_model)
        result = {"status": "complete", "out_dir": str(args.out_dir)}
    else:
        raise SystemExit(f"unknown command: {args.command}")
    print(json.dumps(result, indent=2, ensure_ascii=False, sort_keys=True))


if __name__ == "__main__":
    main()