trainer/validate_dataset.py

from __future__ import annotations

import hashlib
import json
import re
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any

from .choice_guard import assess_choice
from .build_dataset import CATEGORY_COUNTS, REVIEW_STATUS
from .common import APP_ROOT, DATASET_ROOT, atomic_json, ensure_dirs, log_event, update_status


RESPONSE_KEYS = {"speech_choice_id"}
DETERMINISTIC_OUTPUT_KEYS = {
    "next_question_id",
    "workflow_action_id",
    "handoff_requested",
    "cannot_answer",
    "fact_ids_used",
}
SENSITIVE_PATTERNS = [
    re.compile(r"\b\d{3}-\d{2}-\d{4}\b"),
    re.compile(r"\b(?:\d[ -]*?){13,16}\b"),
    re.compile(r"\b[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}\b", re.I),
]
FORBIDDEN_SELECTED_SPEECH = [
    re.compile(pattern, re.I)
    for pattern in (
        r"\b(?:definitely|certainly) (?:safe|an? |have|has)\b",
        r"\b(?:confirms?|proves?) (?:that )?(?:you|the climber|a specific)\b",
        r"\b(?:I (?:will|can)|you should) (?:administer|prescribe|select|take)\b",
        r"\b(?:descend now|start descent|activate emergency)\b",
        r"\b(?:move|start|stop|disable|shutdown|shut down)\b.{0,25}\b(?:robot|motor|joint|actuator)\b",
        r"\bstructured values were received\b",
        r"\bi didn['’]t catch (?:that|it)\b",
        r"\bcould you say it again\b",
    )
]


def read_rows(path: Path) -> list[dict[str, Any]]:
    result: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as handle:
        for line_no, line in enumerate(handle, 1):
            try:
                value = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"{path.name}:{line_no}: invalid JSON: {exc}") from exc
            if not isinstance(value, dict):
                raise ValueError(f"{path.name}:{line_no}: row must be an object")
            result.append(value)
    return result


def expected_policy_flags(request: dict[str, Any]) -> tuple[bool, bool]:
    policy = request.get("policy_state", {})
    review_route = policy.get("review_route")
    answerability = policy.get("answerability")
    return review_route in {"human_review", "operator_review"}, answerability != "supported"


def selected_speech(request: dict[str, Any], response: dict[str, Any]) -> str | None:
    choices = request.get("speech_choices")
    if not isinstance(choices, list):
        return None
    selected = response.get("speech_choice_id")
    for choice in choices:
        if isinstance(choice, dict) and choice.get("id") == selected and isinstance(choice.get("text"), str):
            return choice["text"]
    return None


def validate_contract(
    row_id: str,
    task_type: str,
    request: dict[str, Any],
    response: dict[str, Any],
    questions: dict[str, str],
    workflows: set[str],
    errors: list[str],
) -> None:
    if set(response) != RESPONSE_KEYS:
        errors.append(f"{row_id}: response keys differ from the v5 wording-selection contract")
    if request.get("contract_version") != "sams-wording-selection-v5":
        errors.append(f"{row_id}: wrong request contract version")
    if request.get("task_type") != task_type:
        errors.append(f"{row_id}: task_type mismatch")

    choices = request.get("speech_choices")
    if not isinstance(choices, list) or len(choices) != 4:
        errors.append(f"{row_id}: exactly four speech choices are required")
        choices = []
    choice_ids = [
        choice.get("id") for choice in choices if isinstance(choice, dict) and isinstance(choice.get("id"), str)
    ]
    choice_texts = [
        choice.get("text") for choice in choices if isinstance(choice, dict) and isinstance(choice.get("text"), str)
    ]
    if len(choice_ids) != 4 or len(set(choice_ids)) != 4:
        errors.append(f"{row_id}: speech choice IDs must be four unique strings")
    if len(choice_texts) != 4 or len(set(choice_texts)) != 4:
        errors.append(f"{row_id}: speech choice texts must be four unique strings")
    approved_ids = request.get("approved_speech_choice_ids")
    if not isinstance(approved_ids, list) or set(approved_ids) != set(choice_ids) or len(approved_ids) != 4:
        errors.append(f"{row_id}: every supplied speech choice must be explicitly prevalidated")
    for choice in choices:
        if not isinstance(choice, dict) or not isinstance(choice.get("id"), str):
            continue
        text = choice.get("text")
        if not isinstance(text, str) or len(text) > 300:
            errors.append(f"{row_id}: every supplied speech choice must be text of at most 300 characters")
            continue
        for pattern in FORBIDDEN_SELECTED_SPEECH:
            if pattern.search(text):
                errors.append(f"{row_id}: supplied speech choice contains forbidden pattern {pattern.pattern!r}")
        decision = assess_choice(request, choice["id"])
        if not decision.allowed:
            errors.append(f"{row_id}: supplied speech choice failed guard: {decision.reasons}")
    selected = selected_speech(request, response)
    if selected is None:
        errors.append(f"{row_id}: selected speech choice was not supplied")
    else:
        if len(selected) > 300:
            errors.append(f"{row_id}: selected speech exceeds 300 characters")
        for pattern in FORBIDDEN_SELECTED_SPEECH:
            if pattern.search(selected):
                errors.append(f"{row_id}: selected speech contains forbidden pattern {pattern.pattern!r}")

    allowed_questions = request.get("allowed_question_ids")
    if not isinstance(allowed_questions, list) or not all(isinstance(item, str) for item in allowed_questions):
        errors.append(f"{row_id}: allowed_question_ids must be a string list")
        allowed_questions = []
    if len(allowed_questions) != len(set(allowed_questions)):
        errors.append(f"{row_id}: duplicate allowed question")
    if any(question_id not in questions for question_id in allowed_questions):
        errors.append(f"{row_id}: unknown allowed question")
    expected_question_text = {question_id: questions[question_id] for question_id in allowed_questions if question_id in questions}
    if request.get("question_text_by_id") != expected_question_text:
        errors.append(f"{row_id}: question text map does not exactly match the catalog")
    deterministic = request.get("deterministic_output")
    if not isinstance(deterministic, dict) or set(deterministic) != DETERMINISTIC_OUTPUT_KEYS:
        errors.append(f"{row_id}: deterministic_output keys differ from the v5 contract")
        deterministic = {}
    question = deterministic.get("next_question_id")
    if question is not None and question not in allowed_questions:
        errors.append(f"{row_id}: selected question was not supplied")

    allowed_workflows = request.get("allowed_workflow_action_ids")
    if not isinstance(allowed_workflows, list) or not all(isinstance(item, str) for item in allowed_workflows):
        errors.append(f"{row_id}: allowed_workflow_action_ids must be a string list")
        allowed_workflows = []
    if len(allowed_workflows) != len(set(allowed_workflows)):
        errors.append(f"{row_id}: duplicate allowed workflow action")
    if any(action not in workflows for action in allowed_workflows):
        errors.append(f"{row_id}: unknown allowed workflow action")
    if deterministic.get("workflow_action_id") not in allowed_workflows:
        errors.append(f"{row_id}: selected workflow action was not supplied")

    required_handoff, required_cannot = expected_policy_flags(request)
    if deterministic.get("handoff_requested") is not required_handoff:
        errors.append(f"{row_id}: handoff does not follow deterministic policy state")
    if deterministic.get("cannot_answer") is not required_cannot:
        errors.append(f"{row_id}: cannot_answer does not follow answerability state")

    facts = request.get("facts")
    if not isinstance(facts, list) or not facts:
        errors.append(f"{row_id}: facts must be a non-empty list")
        facts = []
    trusted_ids = {
        fact.get("id")
        for fact in facts
        if isinstance(fact, dict) and fact.get("trust") == "trusted" and isinstance(fact.get("id"), str)
    }
    fact_ids = deterministic.get("fact_ids_used")
    if not isinstance(fact_ids, list) or not all(isinstance(item, str) for item in fact_ids):
        errors.append(f"{row_id}: fact_ids_used must be a string list")
    elif len(fact_ids) != len(set(fact_ids)):
        errors.append(f"{row_id}: fact_ids_used contains duplicates")
    elif not set(fact_ids).issubset(trusted_ids):
        errors.append(f"{row_id}: response cites an absent or untrusted fact")
    elif not fact_ids:
        errors.append(f"{row_id}: response must cite at least one trusted fact")


def validate() -> dict[str, Any]:
    errors: list[str] = []
    warnings: list[str] = []
    manifest_path = DATASET_ROOT / "manifest.json"
    if not manifest_path.is_file():
        return {"ok": False, "errors": ["Missing manifest.json"], "warnings": []}
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    question_payload = json.loads((APP_ROOT / "policy" / "question_catalog.json").read_text(encoding="utf-8"))
    workflow_payload = json.loads((APP_ROOT / "policy" / "workflow_catalog.json").read_text(encoding="utf-8"))
    fixed_warning_payload = json.loads((APP_ROOT / "policy" / "fixed_warning_catalog.json").read_text(encoding="utf-8"))
    workflows = set(workflow_payload["actions"])
    source_ids = {item["id"] for item in manifest["source_manifest"]["sources"]}
    if fixed_warning_payload.get("bypass_llm") is not True:
        errors.append("Fixed warning catalog must bypass the LLM")
    if workflow_payload.get("controller_validation_required") is not True:
        errors.append("Workflow catalog must require controller validation")

    ids: set[str] = set()
    conversation_hashes: set[str] = set()
    family_splits: defaultdict[str, set[str]] = defaultdict(set)
    all_categories: Counter[str] = Counter()
    split_counts: Counter[str] = Counter()
    task_counts: Counter[str] = Counter()
    correct_positions: Counter[int] = Counter()
    selected_speech_texts: set[str] = set()
    training_family_ids: set[str] = set()

    for split in ("train", "validation", "development_test"):
        path = DATASET_ROOT / f"{split}.jsonl"
        if not path.is_file():
            errors.append(f"Missing {path.name}")
            continue
        try:
            rows = read_rows(path)
        except ValueError as exc:
            errors.append(str(exc))
            continue
        for row in rows:
            row_id = str(row.get("id"))
            if row_id in ids:
                errors.append(f"Duplicate ID {row_id}")
            ids.add(row_id)
            split_counts[split] += 1
            category = str(row.get("category"))
            task_type = str(row.get("task_type"))
            all_categories[category] += 1
            task_counts[task_type] += 1
            if row.get("split") != split:
                errors.append(f"{row_id}: split field mismatch")
            if row.get("review_status") != REVIEW_STATUS:
                errors.append(f"{row_id}: unsupported review status")
            if not isinstance(row.get("policy_tags"), list) or not row["policy_tags"]:
                errors.append(f"{row_id}: policy_tags must be non-empty")
            refs = row.get("source_refs")
            if not isinstance(refs, list) or not refs or not set(refs).issubset(source_ids):
                errors.append(f"{row_id}: source_refs are missing or unknown")
            family = str(row.get("scenario_family_id"))
            family_splits[family].add(split)
            training_family_ids.add(family)

            messages = row.get("messages")
            if not isinstance(messages, list) or len(messages) != 3:
                errors.append(f"{row_id}: messages must be exactly system, user, assistant")
                continue
            if [message.get("role") for message in messages if isinstance(message, dict)] != [
                "system",
                "user",
                "assistant",
            ]:
                errors.append(f"{row_id}: invalid role sequence")
                continue
            if any(not isinstance(message.get("content"), str) or not message["content"] for message in messages):
                errors.append(f"{row_id}: empty message content")
                continue
            digest = hashlib.sha256("\n".join(message["content"] for message in messages).encode("utf-8")).hexdigest()
            if digest in conversation_hashes:
                errors.append(f"{row_id}: exact duplicate conversation")
            conversation_hashes.add(digest)
            body = "\n".join(message["content"] for message in messages)
            if any(pattern.search(body) for pattern in SENSITIVE_PATTERNS):
                errors.append(f"{row_id}: possible sensitive identifier")
            try:
                request = json.loads(messages[1]["content"])
                response = json.loads(messages[2]["content"])
            except json.JSONDecodeError:
                errors.append(f"{row_id}: user and assistant messages must be JSON")
                continue
            questions = question_payload["human" if task_type == "climber_checkin" else "robot"]
            validate_contract(row_id, task_type, request, response, questions, workflows, errors)
            choice_ids = [choice["id"] for choice in request.get("speech_choices", []) if isinstance(choice, dict)]
            if response.get("speech_choice_id") in choice_ids:
                correct_positions[choice_ids.index(response["speech_choice_id"])] += 1
            speech = selected_speech(request, response)
            if speech:
                selected_speech_texts.add(speech)

    leaks = {family: sorted(splits) for family, splits in family_splits.items() if len(splits) > 1}
    if leaks:
        errors.append(f"Scenario-family leakage across training splits: {list(leaks.items())[:5]}")
    if dict(all_categories) != CATEGORY_COUNTS:
        errors.append(f"Category counts differ from contract: {dict(all_categories)}")
    if sum(split_counts.values()) != sum(CATEGORY_COUNTS.values()):
        errors.append("Total training/development row count is incorrect")
    if min(split_counts.values(), default=0) < 800:
        errors.append(f"Each split must contain at least 800 rows: {dict(split_counts)}")
    if task_counts["robot_status"] < 1500 or task_counts["climber_checkin"] < 9000:
        errors.append(f"Task balance is too small for one task: {dict(task_counts)}")
    total_positions = sum(correct_positions.values())
    for position in range(4):
        rate = correct_positions[position] / max(1, total_positions)
        if not 0.20 <= rate <= 0.30:
            errors.append(f"Correct speech position {position + 1} is imbalanced at {rate:.3f}")
    if len(selected_speech_texts) < 5000:
        errors.append(f"Selected-speech lexical variety is too low: {len(selected_speech_texts)} unique texts")

    release_path = DATASET_ROOT / "release_gate.jsonl"
    if not release_path.is_file():
        errors.append("Missing release_gate.jsonl")
        release_rows: list[dict[str, Any]] = []
    else:
        release_rows = read_rows(release_path)
        actual_hash = hashlib.sha256(release_path.read_bytes()).hexdigest()
        if actual_hash != manifest.get("release_gate_sha256"):
            errors.append("Release gate SHA-256 does not match the frozen manifest")
        if len(release_rows) != 40:
            errors.append(f"Release gate must contain exactly 40 cases, found {len(release_rows)}")
        release_ids: set[str] = set()
        release_families: set[str] = set()
        for row in release_rows:
            row_id = str(row.get("id"))
            if row_id in release_ids or row_id in ids:
                errors.append(f"Release gate duplicate or leaked ID {row_id}")
            release_ids.add(row_id)
            family = str(row.get("scenario_family_id"))
            if family in release_families or family in training_family_ids:
                errors.append(f"Release gate family is duplicate or leaked: {family}")
            release_families.add(family)
            if row.get("split") != "release_gate":
                errors.append(f"{row_id}: release split marker is wrong")
            messages = row.get("messages", [])
            if len(messages) == 3:
                try:
                    request = json.loads(messages[1]["content"])
                    response = json.loads(messages[2]["content"])
                    questions = question_payload[
                        "human" if row.get("task_type") == "climber_checkin" else "robot"
                    ]
                    validate_contract(
                        row_id,
                        str(row.get("task_type")),
                        request,
                        response,
                        questions,
                        workflows,
                        errors,
                    )
                except (json.JSONDecodeError, KeyError, TypeError) as exc:
                    errors.append(f"{row_id}: invalid release gate record: {exc}")

    required_manifest_flags = {
        "raw_waveforms_fed_to_llm": False,
        "healthbench_used_for_training": False,
        "medsafetybench_used_for_training": False,
        "fixed_warnings_bypass_llm": True,
        "low_confidence_stt_bypasses_llm": True,
        "safety_decisions_bypass_llm": True,
    }
    for key, expected in required_manifest_flags.items():
        if manifest.get(key) is not expected:
            errors.append(f"Manifest flag {key} must be {expected}")
    if manifest.get("review_status") != REVIEW_STATUS:
        errors.append("Manifest review status is incorrect")

    return {
        "ok": not errors,
        "errors": errors,
        "warnings": warnings,
        "split_counts": dict(split_counts),
        "category_counts": dict(sorted(all_categories.items())),
        "task_counts": dict(task_counts),
        "unique_families": len(family_splits),
        "unique_conversations": len(conversation_hashes),
        "unique_selected_speech_texts": len(selected_speech_texts),
        "correct_speech_choice_positions": dict(sorted(correct_positions.items())),
        "release_gate_examples": len(release_rows),
        "release_gate_sha256": manifest.get("release_gate_sha256"),
        "response_contract": "sams-wording-selection-v5",
        "review_status": REVIEW_STATUS,
    }


def main() -> None:
    ensure_dirs()
    update_status("validate_dataset", "Validating v7 approved wording, grounding, split isolation, and release gate")
    report = validate()
    atomic_json(DATASET_ROOT / "validation_report.json", report)
    log_event("dataset_validated", **report)
    if not report["ok"]:
        raise RuntimeError("Dataset validation failed: " + "; ".join(report["errors"][:12]))
    update_status("dataset_valid", "SAMS v7 approved-wording dataset validation passed", validation=report)


if __name__ == "__main__":
    main()