#!/usr/bin/env python3
"""Fail-closed contract for purpose-classifier human review artifacts."""

from __future__ import annotations

import csv
import hashlib
import json
import math
from collections import Counter, defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Sequence

from purpose_data import (
    LABEL_SET,
    SLICES,
    DataError,
    SourceRecord,
    file_sha256,
    normalize_prompt,
    prompt_hash,
    validate_source_record,
    write_json,
)


REVIEW_CSV_FIELDS = (
    "promptHash",
    "source",
    "line",
    "prompt",
    "generatedPurpose",
    "generatedSecondary",
    "generatedMixed",
    "generatedDifficulty",
    "generatedSlice",
    "generatedLanguage",
    "reviewedPurpose",
    "reviewedSecondary",
    "reviewedDifficulty",
    "reviewedSlice",
    "reviewStatus",
    "reviewNotes",
)
REVIEW_STATUSES = frozenset({"accept", "relabel", "reject"})
NONE_SENTINELS = frozenset({"<none>", "none", "null"})
CHANGE_FIELDS = frozenset({"purpose", "secondary", "difficulty", "slice"})
DEFAULT_SAMPLE_SEED = 0xA11D17
DEFAULT_SAMPLE_FRACTION = 0.10


@dataclass(frozen=True)
class ReviewProgress:
    records: int
    accepted: int
    relabeled: int
    rejected: int
    incomplete: int
    decisions: list[dict[str, Any]]

    @property
    def completed(self) -> int:
        return self.records - self.incomplete


@dataclass(frozen=True)
class HumanReviewResult:
    records: list[SourceRecord]
    summary: dict[str, Any]


def _stable_rank(seed: int, record: SourceRecord) -> str:
    material = f"{seed}\0{prompt_hash(record.value['prompt'])}".encode("utf-8")
    return hashlib.sha256(material).hexdigest()


def _review_stratum(record: SourceRecord) -> tuple[str, str, str]:
    language = record.value["lang"].split("-", 1)[0].casefold()
    return record.value["purpose"], record.value["slice"], language


def stratified_review_sample(
    records: Sequence[SourceRecord],
    *,
    fraction: float,
    seed: int,
) -> list[SourceRecord]:
    """Choose exactly round(N*fraction), apportioned by purpose/slice/language."""

    if not records:
        raise DataError("cannot sample an empty review population")
    if not 0.0 < fraction <= 1.0:
        raise DataError("review fraction must be in (0, 1]")

    target = round(len(records) * fraction)
    groups: dict[tuple[str, str, str], list[SourceRecord]] = defaultdict(list)
    for record in records:
        groups[_review_stratum(record)].append(record)

    allocations: dict[tuple[str, str, str], int] = {}
    remainders: list[tuple[float, str, tuple[str, str, str]]] = []
    allocated = 0
    for key in sorted(groups):
        quota = len(groups[key]) * target / len(records)
        base = math.floor(quota)
        allocations[key] = base
        allocated += base
        tie_break = hashlib.sha256(f"{seed}\0{key}".encode("utf-8")).hexdigest()
        remainders.append((quota - base, tie_break, key))
    for _, _, key in sorted(remainders, reverse=True)[: target - allocated]:
        allocations[key] += 1

    selected: list[SourceRecord] = []
    for key in sorted(groups):
        ordered = sorted(groups[key], key=lambda record: _stable_rank(seed, record))
        selected.extend(ordered[: allocations[key]])
    return sorted(selected, key=lambda record: _stable_rank(seed + 1, record))


def _review_row(
    record: SourceRecord,
    source_formatter: Callable[[Path], str],
) -> dict[str, str]:
    value = record.value
    return {
        "promptHash": prompt_hash(value["prompt"]),
        "source": source_formatter(record.source),
        "line": str(record.line),
        "prompt": normalize_prompt(value["prompt"]),
        "generatedPurpose": value["purpose"],
        "generatedSecondary": value["secondary"] or "",
        "generatedMixed": str(value["mixed"]).lower(),
        "generatedDifficulty": str(value["difficulty"]),
        "generatedSlice": value["slice"],
        "generatedLanguage": value["lang"],
        "reviewedPurpose": "",
        "reviewedSecondary": "",
        "reviewedDifficulty": "",
        "reviewedSlice": "",
        "reviewStatus": "",
        "reviewNotes": "",
    }


def write_review_csv(
    path: Path,
    records: Sequence[SourceRecord],
    *,
    source_formatter: Callable[[Path], str] = str,
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=REVIEW_CSV_FIELDS)
        writer.writeheader()
        for record in records:
            writer.writerow(_review_row(record, source_formatter))


def _parse_review_changes(
    row: dict[str, str],
    record: SourceRecord,
    *,
    row_number: int,
) -> dict[str, Any]:
    value = record.value
    changes: dict[str, Any] = {}

    purpose = row["reviewedPurpose"].strip()
    if purpose:
        if purpose not in LABEL_SET:
            raise DataError(f"review CSV row {row_number}: invalid reviewed purpose")
        if purpose != value["purpose"]:
            changes["purpose"] = purpose

    secondary_text = row["reviewedSecondary"].strip()
    secondary = value["secondary"]
    if secondary_text:
        if secondary_text.casefold() in NONE_SENTINELS:
            secondary = None
        elif secondary_text in LABEL_SET:
            secondary = secondary_text
        else:
            raise DataError(f"review CSV row {row_number}: invalid reviewed secondary")
        if secondary != value["secondary"]:
            changes["secondary"] = secondary

    difficulty_text = row["reviewedDifficulty"].strip()
    if difficulty_text:
        try:
            difficulty = float(difficulty_text)
        except ValueError as exc:
            raise DataError(
                f"review CSV row {row_number}: invalid reviewed difficulty"
            ) from exc
        if not math.isfinite(difficulty) or not 0.0 <= difficulty <= 1.0:
            raise DataError(
                f"review CSV row {row_number}: reviewed difficulty must be in [0, 1]"
            )
        if difficulty != float(value["difficulty"]):
            changes["difficulty"] = difficulty

    slice_name = row["reviewedSlice"].strip()
    if slice_name:
        if slice_name not in SLICES:
            raise DataError(f"review CSV row {row_number}: invalid reviewed slice")
        if slice_name != value["slice"]:
            changes["slice"] = slice_name

    candidate = dict(value)
    candidate.update(changes)
    candidate["mixed"] = candidate["secondary"] is not None
    validate_source_record(candidate, f"review CSV row {row_number}")
    return changes


def inspect_review_csv(
    path: Path,
    expected_sample: Sequence[SourceRecord],
    *,
    source_formatter: Callable[[Path], str] = str,
) -> ReviewProgress:
    try:
        with path.open("r", encoding="utf-8", newline="") as handle:
            reader = csv.DictReader(handle)
            if tuple(reader.fieldnames or ()) != REVIEW_CSV_FIELDS:
                raise DataError(
                    f"{path}: review CSV fields do not match the current schema; "
                    "regenerate the blank artifact before reviewing"
                )
            rows = list(reader)
    except (OSError, UnicodeError, csv.Error) as exc:
        raise DataError(f"{path}: cannot read review CSV: {exc}") from exc

    expected_by_hash = {
        prompt_hash(record.value["prompt"]): record for record in expected_sample
    }
    if len(expected_by_hash) != len(expected_sample):
        raise DataError("human-review sample contains duplicate prompt hashes")
    if len(rows) != len(expected_sample):
        raise DataError(
            f"{path}: expected {len(expected_sample)} review rows, found {len(rows)}"
        )

    seen: set[str] = set()
    decisions: list[dict[str, Any]] = []
    counts: Counter[str] = Counter()
    incomplete = 0
    for row_number, row in enumerate(rows, 2):
        if None in row or any(row.get(field) is None for field in REVIEW_CSV_FIELDS):
            raise DataError(f"{path}:{row_number}: malformed review CSV column count")
        row_hash = row["promptHash"].strip()
        if row_hash in seen:
            raise DataError(f"{path}:{row_number}: duplicate promptHash {row_hash}")
        seen.add(row_hash)
        record = expected_by_hash.get(row_hash)
        if record is None:
            raise DataError(
                f"{path}:{row_number}: promptHash is not in the deterministic sample"
            )

        expected = _review_row(record, source_formatter)
        for field in REVIEW_CSV_FIELDS[:10]:
            actual = row[field]
            if field == "generatedDifficulty":
                try:
                    matches = float(actual) == float(expected[field])
                except ValueError:
                    matches = False
            else:
                matches = actual == expected[field]
            if not matches:
                raise DataError(
                    f"{path}:{row_number}: generated field {field} no longer "
                    "matches the curated corpus"
                )

        status = row["reviewStatus"].strip().casefold()
        reviewer_values = [row[field].strip() for field in REVIEW_CSV_FIELDS[10:14]]
        notes = row["reviewNotes"].strip()
        if not status:
            if any(reviewer_values) or notes:
                raise DataError(
                    f"{path}:{row_number}: reviewer fields require a reviewStatus"
                )
            incomplete += 1
            continue
        if status not in REVIEW_STATUSES:
            raise DataError(f"{path}:{row_number}: invalid reviewStatus {status!r}")

        changes = _parse_review_changes(row, record, row_number=row_number)
        if status == "accept":
            if any(reviewer_values):
                raise DataError(
                    f"{path}:{row_number}: accept must leave reviewed fields blank"
                )
        elif status == "reject":
            if any(reviewer_values):
                raise DataError(
                    f"{path}:{row_number}: reject must leave reviewed fields blank"
                )
            if not notes:
                raise DataError(f"{path}:{row_number}: reject requires reviewNotes")
        else:
            if not changes:
                raise DataError(
                    f"{path}:{row_number}: relabel must change at least one field"
                )
            if not notes:
                raise DataError(f"{path}:{row_number}: relabel requires reviewNotes")

        decision: dict[str, Any] = {
            "promptHash": row_hash,
            "status": status,
        }
        if changes:
            decision["changes"] = changes
        if notes:
            decision["notes"] = notes
        decisions.append(decision)
        counts[status] += 1

    missing = sorted(set(expected_by_hash) - seen)
    if missing:
        raise DataError(f"{path}: deterministic sample rows are missing")
    return ReviewProgress(
        records=len(rows),
        accepted=counts["accept"],
        relabeled=counts["relabel"],
        rejected=counts["reject"],
        incomplete=incomplete,
        decisions=decisions,
    )


def finalize_human_review(
    csv_path: Path,
    output_path: Path,
    population: Sequence[SourceRecord],
    *,
    dataset_version: str,
    fraction: float,
    seed: int,
    source_formatter: Callable[[Path], str] = str,
) -> dict[str, Any]:
    sample = stratified_review_sample(population, fraction=fraction, seed=seed)
    progress = inspect_review_csv(
        csv_path,
        sample,
        source_formatter=source_formatter,
    )
    if progress.incomplete:
        raise DataError(
            f"{csv_path}: human review is incomplete: {progress.completed}/"
            f"{progress.records} rows completed"
        )

    summary = {
        "accepted": progress.accepted,
        "relabeled": progress.relabeled,
        "rejected": progress.rejected,
    }
    artifact = {
        "schemaVersion": 1,
        "datasetVersion": dataset_version,
        "status": "complete",
        "populationRecords": len(population),
        "sampleFraction": fraction,
        "sampleRecords": len(sample),
        "seed": seed,
        "sourceCSVSha256": file_sha256(csv_path),
        "summary": summary,
        "decisions": progress.decisions,
    }
    if output_path.exists():
        try:
            existing = json.loads(output_path.read_text(encoding="utf-8"))
        except (OSError, UnicodeError, json.JSONDecodeError) as exc:
            raise DataError(
                f"{output_path}: cannot verify existing human-review ledger: {exc}"
            ) from exc
        if existing != artifact:
            raise DataError(
                f"{output_path}: refusing to replace a different human-review ledger"
            )
        return artifact
    write_json(output_path, artifact)
    return artifact


def apply_completed_human_review(
    path: Path,
    population: Sequence[SourceRecord],
    *,
    dataset_version: str,
) -> HumanReviewResult:
    try:
        artifact = json.loads(path.read_text(encoding="utf-8"))
        if artifact["schemaVersion"] != 1:
            raise DataError(f"{path}: unsupported human-review schema")
        if artifact["datasetVersion"] != dataset_version:
            raise DataError(f"{path}: human-review dataset version does not match")
        if artifact["status"] != "complete":
            raise DataError(f"{path}: human review is not complete")
        fraction = float(artifact["sampleFraction"])
        seed = int(artifact["seed"])
        decisions = artifact["decisions"]
    except DataError:
        raise
    except (OSError, UnicodeError, json.JSONDecodeError, KeyError, TypeError, ValueError) as exc:
        raise DataError(f"{path}: cannot read completed human review: {exc}") from exc
    if not isinstance(decisions, list):
        raise DataError(f"{path}: human-review decisions must be an array")
    if artifact.get("populationRecords") != len(population):
        raise DataError(f"{path}: reviewed population no longer matches the corpus")

    sample = stratified_review_sample(population, fraction=fraction, seed=seed)
    expected_hashes = {prompt_hash(record.value["prompt"]) for record in sample}
    if artifact.get("sampleRecords") != len(sample):
        raise DataError(f"{path}: reviewed sample size no longer matches the corpus")

    by_hash: dict[str, dict[str, Any]] = {}
    counts: Counter[str] = Counter()
    for index, decision in enumerate(decisions, 1):
        if not isinstance(decision, dict):
            raise DataError(f"{path}: human-review decision {index} must be an object")
        row_hash = decision.get("promptHash")
        status = decision.get("status")
        changes = decision.get("changes", {})
        if not isinstance(row_hash, str) or not isinstance(status, str):
            raise DataError(f"{path}: invalid human-review decision {index}")
        if row_hash in by_hash:
            raise DataError(f"{path}: duplicate human-review promptHash")
        if row_hash not in expected_hashes or status not in REVIEW_STATUSES:
            raise DataError(f"{path}: invalid human-review decision {index}")
        if not isinstance(changes, dict) or not set(changes) <= CHANGE_FIELDS:
            raise DataError(f"{path}: invalid changes in human-review decision {index}")
        if status != "relabel" and changes:
            raise DataError(f"{path}: only relabel decisions may contain changes")
        if status == "relabel" and not changes:
            raise DataError(f"{path}: relabel decision {index} has no changes")
        if "purpose" in changes and (
            not isinstance(changes["purpose"], str)
            or changes["purpose"] not in LABEL_SET
        ):
            raise DataError(f"{path}: invalid purpose in human-review decision {index}")
        if "secondary" in changes and (
            changes["secondary"] is not None
            and (
                not isinstance(changes["secondary"], str)
                or changes["secondary"] not in LABEL_SET
            )
        ):
            raise DataError(f"{path}: invalid secondary in human-review decision {index}")
        if "slice" in changes and (
            not isinstance(changes["slice"], str) or changes["slice"] not in SLICES
        ):
            raise DataError(f"{path}: invalid slice in human-review decision {index}")
        by_hash[row_hash] = decision
        counts[status] += 1
    if set(by_hash) != expected_hashes:
        raise DataError(f"{path}: decisions do not exactly cover the deterministic sample")

    expected_summary = {
        "accepted": counts["accept"],
        "relabeled": counts["relabel"],
        "rejected": counts["reject"],
    }
    if artifact.get("summary") != expected_summary:
        raise DataError(f"{path}: human-review summary does not match its decisions")

    result: list[SourceRecord] = []
    for record in population:
        decision = by_hash.get(prompt_hash(record.value["prompt"]))
        if decision is None or decision["status"] == "accept":
            result.append(record)
            continue
        if decision["status"] == "reject":
            continue
        value = dict(record.value)
        value.update(decision["changes"])
        value["mixed"] = value["secondary"] is not None
        validate_source_record(value, f"{path}: {decision['promptHash']}")
        result.append(SourceRecord(value=value, source=record.source, line=record.line))

    return HumanReviewResult(
        records=result,
        summary={
            **expected_summary,
            "sampleRecords": len(sample),
            "retainedRecords": len(result),
        },
    )