"""Bounded illustrative assistant-only SFT data audit; does not train a model."""
from dataclasses import asdict, dataclass, fields
import hashlib
import json


def text(value):
    if type(value) is not str or not value.strip() or len(value) > 4096:
        raise ValueError("nonempty bounded string required")
    return value


def count(value, minimum=1):
    if type(value) is not int or not minimum <= value <= 1_000_000:
        raise ValueError("bounded integer required")


def digest(value):
    return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":"),
                                     allow_nan=False).encode()).hexdigest()


@dataclass(frozen=True)
class Turn:
    role: str
    tokens: tuple[int, ...]
    loss_mask: tuple[int, ...]

    def __post_init__(self):
        text(self.role)
        if self.role not in ("system", "user", "assistant"):
            raise ValueError("unknown role")
        for name in ("tokens", "loss_mask"):
            value = getattr(self, name)
            if type(value) not in (list, tuple) or not 1 <= len(value) <= 4096:
                raise ValueError("bounded token sequence required")
            object.__setattr__(self, name, tuple(value))
        for token in self.tokens:
            count(token, 0)
        if len(self.tokens) != len(self.loss_mask):
            raise ValueError("mask length mismatch")
        expected = 1 if self.role == "assistant" else 0
        if any(type(mask) is not int or mask != expected for mask in self.loss_mask):
            raise ValueError("assistant-only objective required")


@dataclass(frozen=True)
class Row:
    row_id: str
    conversation_id: str
    source_group: str
    split: str
    split_provenance: str
    license: str
    pii_review: str
    turns: tuple[Turn, ...]

    def __post_init__(self):
        for name in ("row_id", "conversation_id", "source_group", "split_provenance", "license", "split", "pii_review"):
            text(getattr(self, name))
        if self.split not in ("train", "eval") or self.pii_review != "approved-redacted":
            raise ValueError("split and completed PII review required")
        if type(self.turns) not in (tuple, list) or not 2 <= len(self.turns) <= 64:
            raise ValueError("bounded complete conversation required")
        copied = []
        for item in self.turns:
            if type(item) is not Turn:
                raise ValueError("concrete Turn required")
            copied.append(Turn(item.role, item.tokens, item.loss_mask))
        object.__setattr__(self, "turns", tuple(copied))
        roles = [item.role for item in copied]
        if roles[0] == "system":
            roles = roles[1:]
        if len(roles) < 2 or len(roles) % 2 or any(role != ("user" if i % 2 == 0 else "assistant") for i, role in enumerate(roles)):
            raise ValueError("conversation must alternate user/assistant and end at assistant boundary")


@dataclass(frozen=True)
class DataContract:
    scope: str
    model: str
    tokenizer: str
    data_version: str
    policy: str
    provenance: str
    dedup_version: str
    eval_protocol: str
    eval_content_digest: str
    allowed_licenses: tuple[str, ...]
    min_eval_rows: int
    max_tokens: int
    claim: str
    rows: tuple[Row, ...]

    def __post_init__(self):
        for name in ("scope", "model", "tokenizer", "data_version", "policy", "provenance", "dedup_version", "eval_protocol", "claim"):
            text(getattr(self, name))
        if type(self.eval_content_digest) is not str or len(self.eval_content_digest) != 64 or any(c not in "0123456789abcdef" for c in self.eval_content_digest):
            raise ValueError("evaluation content SHA-256 required")
        count(self.min_eval_rows)
        count(self.max_tokens)
        if self.claim != "DATA_CONTRACT_ONLY_NO_GENERALIZATION_CLAIM":
            raise ValueError("unqualified memorization/generalization claim")
        if type(self.allowed_licenses) not in (list, tuple) or not 1 <= len(self.allowed_licenses) <= 100:
            raise ValueError("license allowlist required")
        object.__setattr__(self, "allowed_licenses", tuple(text(x) for x in self.allowed_licenses))
        if len(set(self.allowed_licenses)) != len(self.allowed_licenses):
            raise ValueError("duplicate license")
        if type(self.rows) not in (list, tuple) or not 2 <= len(self.rows) <= 10000:
            raise ValueError("bounded train and evaluation rows required")
        copied = []
        for row in self.rows:
            if type(row) is not Row:
                raise ValueError("concrete Row required")
            copied.append(Row(**{f.name: getattr(row, f.name) for f in fields(Row)}))
        object.__setattr__(self, "rows", tuple(copied))
        for name in ("row_id", "conversation_id"):
            if len({getattr(row, name) for row in copied}) != len(copied):
                raise ValueError("duplicate " + name)
        seen_contents = set()
        source_splits = {}
        for row in copied:
            if row.license not in self.allowed_licenses:
                raise ValueError("license not allowed")
            if sum(len(turn.tokens) for turn in row.turns) > self.max_tokens:
                raise ValueError("conversation exceeds boundary; truncation forbidden")
            content = digest([{"role": turn.role, "tokens": turn.tokens} for turn in row.turns])
            if content in seen_contents:
                raise ValueError("duplicate content or split leakage")
            seen_contents.add(content)
            previous = source_splits.setdefault(row.source_group, row.split)
            if previous != row.split:
                raise ValueError("source-group split leakage")
        if not any(row.split == "train" for row in copied) or sum(row.split == "eval" for row in copied) < self.min_eval_rows:
            raise ValueError("unqualified evaluation split")
        eval_contents = [asdict(row) for row in copied if row.split == "eval"]
        if digest(eval_contents) != self.eval_content_digest:
            raise ValueError("evaluation contents changed")


@dataclass(frozen=True)
class DataAudit:
    contract_digest: str
    train_rows: int
    eval_rows: int
    supervised_tokens: int
    claim: str


def audit(contract):
    if type(contract) is not DataContract:
        raise ValueError("concrete DataContract required")
    contract = DataContract(**{f.name: getattr(contract, f.name) for f in fields(DataContract)})
    return DataAudit(digest(asdict(contract)), sum(row.split == "train" for row in contract.rows),
                     sum(row.split == "eval" for row in contract.rows),
                     sum(sum(turn.loss_mask) for row in contract.rows if row.split == "train" for turn in row.turns), contract.claim)


def example_contract():
    rows = tuple(Row("row-" + split, "conversation-" + split, "source-" + split, split,
                     "group holdout before authoring; synthetic fixture", "CC0", "approved-redacted",
                     (Turn("user", (seed, seed + 1), (0, 0)), Turn("assistant", (seed + 2, seed + 3), (1, 1))))
                 for split, seed in (("train", 10), ("eval", 20)))
    return DataContract("illustrative-support", "base-v1", "tokenizer-v1", "data-v1", "assistant-only-v1",
                        "synthetic; no customer data", "exact-token-v1", "held-out rubric-v1",
                        digest([asdict(rows[1])]), ("CC0",), 1, 128,
                        "DATA_CONTRACT_ONLY_NO_GENERALIZATION_CLAIM", rows)


def main():
    result = audit(example_contract())
    print("example=illustrative_only")
    print(f"rows=train:{result.train_rows},eval:{result.eval_rows}")
    print(f"supervised_train_tokens={result.supervised_tokens}")
    print("claim=" + result.claim)


if __name__ == "__main__":
    main()