"""Audit token-span coverage for a declared chunking policy; do not tokenize documents."""
from dataclasses import asdict, dataclass, fields
import hashlib
import json
import math
import sys


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


def sha256_hex(value):
    if type(value) is not str or len(value) != 64 or any(c not in "0123456789abcdef" for c in value):
        raise ValueError("lowercase SHA-256 required")
    return value


def integer(value, lower, upper):
    if type(value) is not int or not lower <= value <= upper:
        raise ValueError("bounded integer, not bool, required")
    return value


def number(value, lower, upper):
    if type(value) not in (int, float) or not math.isfinite(value):
        raise ValueError("finite real, not bool, required")
    if value != 0 and abs(value) < sys.float_info.min:
        raise ValueError("subnormal rejected")
    if not lower <= value <= upper:
        raise ValueError("number out of range")
    return float(value)


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


def identities(value, label, lower=1, upper=100):
    if type(value) not in (tuple, list) or not lower <= len(value) <= upper:
        raise ValueError(f"bounded {label} sequence required")
    copied = tuple(text(item) for item in value)
    if len(set(copied)) != len(copied):
        raise ValueError(f"duplicate {label}")
    return copied


@dataclass(frozen=True)
class ChunkingPolicy:
    policy_id: str
    strategy: str
    version: str
    tokenizer_revision: str
    boundary_revision: str
    embedding_contract_digest: str
    max_chunk_tokens: int
    overlap_tokens: int
    owner: str
    rollback: str

    def __post_init__(self):
        for name in ("policy_id", "strategy", "version", "tokenizer_revision", "boundary_revision", "owner", "rollback"):
            text(getattr(self, name))
        if self.strategy not in ("fixed-window", "sentence-window", "structure-aware"):
            raise ValueError("unsupported chunking strategy")
        sha256_hex(self.embedding_contract_digest)
        integer(self.max_chunk_tokens, 1, 1_000_000)
        integer(self.overlap_tokens, 0, self.max_chunk_tokens - 1)


@dataclass(frozen=True)
class SourceDocument:
    document_id: str
    revision: str
    content_digest: str
    token_count: int

    def __post_init__(self):
        text(self.document_id)
        text(self.revision)
        sha256_hex(self.content_digest)
        integer(self.token_count, 1, 100_000_000)


@dataclass(frozen=True)
class Chunk:
    chunk_id: str
    document_id: str
    document_revision: str
    start_token: int
    end_token: int
    heading_path: tuple[str, ...]
    content_digest: str

    def __post_init__(self):
        for name in ("chunk_id", "document_id", "document_revision"):
            text(getattr(self, name))
        integer(self.start_token, 0, 100_000_000)
        integer(self.end_token, 1, 100_000_000)
        if self.end_token <= self.start_token:
            raise ValueError("chunk must be a nonempty half-open span")
        object.__setattr__(self, "heading_path", identities(self.heading_path, "heading path", 1, 32))
        sha256_hex(self.content_digest)


@dataclass(frozen=True)
class AnswerCase:
    query_id: str
    question_digest: str
    document_id: str
    document_revision: str
    answer_start: int
    answer_end: int
    retrieved_chunk_ids: tuple[str, ...]

    def __post_init__(self):
        for name in ("query_id", "document_id", "document_revision"):
            text(getattr(self, name))
        sha256_hex(self.question_digest)
        integer(self.answer_start, 0, 100_000_000)
        integer(self.answer_end, 1, 100_000_000)
        if self.answer_end <= self.answer_start:
            raise ValueError("answer must be a nonempty half-open span")
        object.__setattr__(self, "retrieved_chunk_ids", identities(self.retrieved_chunk_ids, "retrieved chunk", 1, 100))


@dataclass(frozen=True)
class ChunkingAuditContract:
    contract_id: str
    scope: str
    version: str
    policy: ChunkingPolicy
    corpus_digest: str
    case_set_digest: str
    top_k: int
    min_cases: int
    min_span_coverage: float
    min_retrieval_coverage: float
    max_mean_minimal_waste_tokens: float
    documents: tuple[SourceDocument, ...]
    chunks: tuple[Chunk, ...]
    cases: tuple[AnswerCase, ...]

    def __post_init__(self):
        for name in ("contract_id", "scope", "version"):
            text(getattr(self, name))
        if type(self.policy) is not ChunkingPolicy:
            raise ValueError("concrete ChunkingPolicy required")
        object.__setattr__(self, "policy", ChunkingPolicy(**{field.name: getattr(self.policy, field.name) for field in fields(ChunkingPolicy)}))
        sha256_hex(self.corpus_digest)
        sha256_hex(self.case_set_digest)
        integer(self.top_k, 1, 100)
        integer(self.min_cases, 1, 100_000)
        number(self.min_span_coverage, 0, 1)
        number(self.min_retrieval_coverage, 0, 1)
        number(self.max_mean_minimal_waste_tokens, 0, 100_000_000)
        copied_groups = []
        for value, cls, label, upper in (
            (self.documents, SourceDocument, "documents", 100_000),
            (self.chunks, Chunk, "chunks", 1_000_000),
            (self.cases, AnswerCase, "cases", 100_000),
        ):
            if type(value) not in (tuple, list) or not 1 <= len(value) <= upper:
                raise ValueError(f"bounded {label} required")
            copied = []
            for item in value:
                if type(item) is not cls:
                    raise ValueError(f"concrete {cls.__name__} required")
                copied.append(cls(**{field.name: getattr(item, field.name) for field in fields(cls)}))
            copied_groups.append(tuple(copied))
        documents, chunks, cases = copied_groups
        object.__setattr__(self, "documents", documents)
        object.__setattr__(self, "chunks", chunks)
        object.__setattr__(self, "cases", cases)
        if len({item.document_id for item in documents}) != len(documents):
            raise ValueError("duplicate document identity")
        if len({item.chunk_id for item in chunks}) != len(chunks):
            raise ValueError("duplicate chunk identity")
        if len({item.query_id for item in cases}) != len(cases):
            raise ValueError("duplicate query identity")
        if digest([asdict(item) for item in documents]) != self.corpus_digest:
            raise ValueError("corpus content mismatch")
        if digest([asdict(item) for item in cases]) != self.case_set_digest:
            raise ValueError("case set content mismatch")
        if len(cases) < self.min_cases:
            raise ValueError("insufficient cases")
        documents_by_id = {item.document_id: item for item in documents}
        chunks_by_id = {item.chunk_id: item for item in chunks}
        for chunk in chunks:
            document = documents_by_id.get(chunk.document_id)
            if document is None or chunk.document_revision != document.revision:
                raise ValueError("chunk source binding mismatch")
            if chunk.end_token > document.token_count:
                raise ValueError("chunk exceeds source document")
            if chunk.end_token - chunk.start_token > self.policy.max_chunk_tokens:
                raise ValueError("chunk exceeds policy size")
        for document in documents:
            source_chunks = [item for item in chunks if item.document_id == document.document_id]
            if not source_chunks or source_chunks != sorted(source_chunks, key=lambda item: (item.start_token, item.end_token, item.chunk_id)):
                raise ValueError("chunks must be present in canonical source order")
            if source_chunks[0].start_token != 0 or source_chunks[-1].end_token != document.token_count:
                raise ValueError("document edges are uncovered")
            for previous, current in zip(source_chunks, source_chunks[1:]):
                if current.start_token > previous.end_token:
                    raise ValueError("chunk gap")
                if previous.end_token - current.start_token > self.policy.overlap_tokens:
                    raise ValueError("overlap exceeds policy")
            if len({(item.start_token, item.end_token) for item in source_chunks}) != len(source_chunks):
                raise ValueError("duplicate chunk span")
        for case in cases:
            document = documents_by_id.get(case.document_id)
            if document is None or case.document_revision != document.revision or case.answer_end > document.token_count:
                raise ValueError("answer source binding mismatch")
            if len(case.retrieved_chunk_ids) != self.top_k or any(item not in chunks_by_id for item in case.retrieved_chunk_ids):
                raise ValueError("retrieval evidence must contain known exact top-k identities")


@dataclass(frozen=True)
class ChunkingAudit:
    contract_digest: str
    span_coverage: float
    retrieval_coverage: float
    mean_minimal_waste_tokens: float
    decision: str
    claim: str


def audit(contract):
    if type(contract) is not ChunkingAuditContract:
        raise ValueError("concrete ChunkingAuditContract required")
    contract = ChunkingAuditContract(**{field.name: getattr(contract, field.name) for field in fields(ChunkingAuditContract)})
    documents = {item.document_id: item for item in contract.documents}
    chunks = {item.chunk_id: item for item in contract.chunks}
    representable = 0
    retrieved = 0
    wastes = []
    for case in contract.cases:
        containers = [item for item in contract.chunks if item.document_id == case.document_id
                      and item.document_revision == case.document_revision
                      and item.start_token <= case.answer_start and item.end_token >= case.answer_end]
        if containers:
            representable += 1
            wastes.append(min(item.end_token - item.start_token for item in containers) - (case.answer_end - case.answer_start))
        else:
            wastes.append(documents[case.document_id].token_count - (case.answer_end - case.answer_start))
        if any(chunks[item].document_id == case.document_id
               and chunks[item].document_revision == case.document_revision
               and chunks[item].start_token <= case.answer_start
               and chunks[item].end_token >= case.answer_end for item in case.retrieved_chunk_ids):
            retrieved += 1
    span_coverage = representable / len(contract.cases)
    retrieval_coverage = retrieved / len(contract.cases)
    mean_waste = sum(wastes) / len(wastes)
    passed = (span_coverage >= contract.min_span_coverage
              and retrieval_coverage >= contract.min_retrieval_coverage
              and mean_waste <= contract.max_mean_minimal_waste_tokens)
    return ChunkingAudit(digest(asdict(contract)), span_coverage, retrieval_coverage, mean_waste,
                         "PASS_DECLARED_COVERAGE_GATES" if passed else "BLOCK",
                         "TOKEN_SPAN_FIXTURE_ONLY")


def example_contract():
    policy = ChunkingPolicy("support-fixed-v1", "fixed-window", "1", "tokenizer-example-001", "window-v1",
                            digest({"embedding-contract": "example-v1"}), 6, 2,
                            "example-search-team", "restore support-fixed-v0")
    documents = (SourceDocument("policy", "7", digest({"tokens": list(range(12))}), 12),)
    chunks = (
        Chunk("policy-0", "policy", "7", 0, 6, ("Returns",), digest({"range": [0, 6]})),
        Chunk("policy-1", "policy", "7", 4, 10, ("Returns", "Timing"), digest({"range": [4, 10]})),
        Chunk("policy-2", "policy", "7", 8, 12, ("Exceptions",), digest({"range": [8, 12]})),
    )
    cases = (
        AnswerCase("refund-window", digest({"question": "refund window"}), "policy", "7", 4, 6,
                   ("policy-0", "policy-1")),
        AnswerCase("exception", digest({"question": "exception"}), "policy", "7", 9, 11,
                   ("policy-2", "policy-1")),
    )
    return ChunkingAuditContract("support-chunking-v1", "illustrative support retrieval", "1", policy,
                                 digest([asdict(item) for item in documents]),
                                 digest([asdict(item) for item in cases]),
                                 2, 2, 1.0, 1.0, 3.0, documents, chunks, cases)


def main():
    result = audit(example_contract())
    print("example=illustrative_only")
    print(f"span_coverage={result.span_coverage:.3f}")
    print(f"retrieval_coverage={result.retrieval_coverage:.3f}")
    print(f"mean_minimal_waste_tokens={result.mean_minimal_waste_tokens:.3f}")
    print("decision=" + result.decision)
    print("claim=" + result.claim)


if __name__ == "__main__":
    main()