"""Strict rows and whole-corpus leakage detection."""
from __future__ import annotations

import hashlib
import json
import re
import unicodedata
from pathlib import Path
from typing import Annotated, Literal

from pydantic import Field, JsonValue, model_validator

from .contracts import DecisionInput, Nonempty, StrictModel
from .render import canonical


class Provenance(StrictModel):
    license: Nonempty
    source_revision: Nonempty
    label_method: Literal["program_oracle", "public_human_label", "original_human_label", "teacher_only", "format_example"]
    human_reviewed: bool
    review_status: Literal["pending", "reviewed", "rejected", "adjudication"]
    oracle_id: Nonempty | None = None
    teacher_model: Nonempty | None = None
    teacher_request_id: Nonempty | None = None
    repository: str | None = None
    original_path: str | None = None
    original_sha256: str | None = None
    original_split: str | None = None
    original_sample_id: str | None = None
    original_label: str | None = None
    attribution: str | None = None
    transformation: str | None = None
    text_modified: bool | None = None
    label_modified: bool | None = None
    authorship: str | None = None
    generator_seed: int | None = None
    rule_version: str | None = None
    counterfactual_ancestry: str | None = None
    ancestry_known: bool | None = None
    oracle_facts: dict[str, JsonValue] | None = None
    reviewer_id: str | None = None
    reviewed_at: str | None = None
    family_audited: bool | None = None
    family_audit_id: str | None = None
    family_audit_sha256: str | None = None
    acceptance_basis: Literal['original_public_label','individual_human_review','family_audited_program'] | None = None
    note: str = ""

    @model_validator(mode="after")
    def methods(self):
        for name in ("license", "source_revision"):
            if not getattr(self, name).strip():
                raise ValueError(f"provenance {name} must not be blank")
        if self.human_reviewed != (self.review_status == "reviewed"):
            raise ValueError("human_reviewed must agree with review_status")
        if self.label_method == "program_oracle" and not self.oracle_id:
            raise ValueError("program_oracle requires an oracle_id")
        if self.label_method == "teacher_only" and not (self.teacher_model and self.teacher_request_id):
            raise ValueError("teacher_only requires teacher model and request identifiers")
        if self.label_method == "original_human_label" and not self.human_reviewed:
            raise ValueError("original human labels require recorded review")
        return self


FORBIDDEN_INPUT_KEYS = {"label", "labels", "gold_label", "teacher_rationale", "teacher_probabilities", "rationale", "target_probs"}


def protect_state(value: object) -> None:
    if isinstance(value, dict):
        bad = {str(k).casefold() for k in value} & FORBIDDEN_INPUT_KEYS
        if bad:
            raise ValueError(f"supervision fields {sorted(bad)} must stay outside state")
        for part in value.values():
            protect_state(part)
    elif isinstance(value, list):
        for part in value:
            protect_state(part)


class DecisionRow(DecisionInput):
    id: Nonempty
    source_id: Nonempty
    source_family: Nonempty
    template_family: Nonempty
    group_id: Nonempty
    language: Literal["en", "hi", "hinglish", "other"]
    split: Literal["train", "validation", "calibration", "sealed_test", "example"]
    label: Annotated[int, Field(ge=0)]
    provenance: Provenance
    teacher_rationale: Annotated[str, Field(max_length=4096)] | None = None
    difficulty_flags: list[str] = Field(default_factory=list)
    curation_status: Literal["candidate", "accepted", "quarantined"] = "candidate"

    @model_validator(mode="after")
    def labels(self):
        for name in ("id", "source_id", "source_family", "template_family", "group_id"):
            if not getattr(self, name).strip():
                raise ValueError(f"{name} must not be blank")
        if self.label >= len(self.options):
            raise ValueError(f"label {self.label} is outside 0..{len(self.options)-1}")
        protect_state(self.state)
        if self.provenance.review_status == "rejected":
            raise ValueError("rejected labels cannot enter a dataset")
        if self.provenance.label_method == "format_example" and self.split != "example":
            raise ValueError("format examples must stay in the example split")
        if self.provenance.label_method == "teacher_only" and self.split != "train":
            raise ValueError("teacher labels are training-only; independent evaluation needs gold labels")
        if self.split == "train" and self.provenance.original_split in {"test", "oos_test"}:
            raise ValueError("official public test data cannot enter training")
        if self.teacher_rationale is not None and self.split != "train":
            raise ValueError("teacher rationale is training audit metadata only")
        if self.curation_status == "accepted" and self.provenance.label_method != "public_human_label":
            if not (self.provenance.human_reviewed and self.provenance.reviewer_id and self.provenance.reviewed_at):
                p=self.provenance
                if not (p.label_method=='program_oracle' and p.family_audited is True and p.family_audit_id and p.family_audit_sha256
                        and p.acceptance_basis=='family_audited_program' and self.language=='en' and not p.label_modified):
                    raise ValueError("acceptance requires original public gold, named dated human review or externally verified family audit")
        if self.curation_status == "accepted" and self.provenance.review_status == "adjudication":
            raise ValueError("unresolved adjudication cannot be accepted")
        return self

    def model_input(self) -> DecisionInput:
        return DecisionInput.model_validate({k: getattr(self, k) for k in ("task_type", "state", "instructions", "options")})


def strict_json(text: str):
    def pairs(items):
        result = {}
        for key, value in items:
            if key in result:
                raise ValueError(f"duplicate JSON object key: {key}")
            result[key] = value
        return result

    def nonfinite(value):
        raise ValueError(f"nonfinite JSON literal: {value}")
    return json.loads(text, object_pairs_hook=pairs, parse_constant=nonfinite)


def normalized(value: object):
    if isinstance(value, str):
        return re.sub(r"\s+", " ", unicodedata.normalize("NFKC", value).casefold()).strip()
    if isinstance(value, dict):
        return {k: normalized(v) for k, v in value.items()}
    if isinstance(value, list):
        return [normalized(v) for v in value]
    return value


def fingerprint(row: DecisionRow, *, normalize: bool = False) -> str:
    payload = row.model_input().model_dump()
    if normalize:
        payload = normalized(payload)
    # Choice permutations do not create independent data; score order stays significant.
    if row.task_type == "choice":
        payload["options"] = sorted(payload["options"], key=lambda o: o["key"])
    return hashlib.sha256(canonical(payload).encode("utf-8")).hexdigest()


def validate_corpus(rows: list[DecisionRow]) -> None:
    if not rows:
        raise ValueError("dataset is empty")
    ids, exact, similar = {}, {}, {}
    lineage = {field: {} for field in ("source_family", "template_family", "group_id")}
    for row in rows:
        if row.id in ids:
            raise ValueError(f"{row.id}: duplicate ID; first seen at row {ids[row.id]}")
        ids[row.id] = len(ids) + 1
        for field, seen in lineage.items():
            key = getattr(row, field)
            if key in seen and seen[key] != row.split:
                raise ValueError(f"{row.id}: split leakage: {field}={key} crosses {seen[key]} and {row.split}")
            seen[key] = row.split
        for normalize, seen in ((False, exact), (True, similar)):
            key = fingerprint(row, normalize=normalize)
            if key in seen:
                kind = "normalized" if normalize else "exact"
                raise ValueError(f"{row.id}: {kind} duplicate model input of {seen[key]}")
            seen[key] = row.id


def load_rows(paths: list[Path]) -> list[DecisionRow]:
    rows = []
    for path in paths:
        for line_number, line in enumerate(path.read_text(encoding="utf-8-sig").splitlines(), 1):
            try:
                if not line.strip():
                    raise ValueError("blank JSONL row")
                rows.append(DecisionRow.model_validate(strict_json(line)))
            except (ValueError, TypeError) as exc:
                raise ValueError(f"{path}:{line_number}: {exc}") from exc
    validate_corpus(rows)
    return rows