"""Deterministic distributed-training topology audit over invented evidence."""

from collections import Counter
from dataclasses import asdict, dataclass, fields
import hashlib
import json
import math
import re


IDENTITY = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.:/@-]{0,127}")
SHA256 = re.compile(r"[a-f0-9]{64}")


def identity(value):
    if type(value) is not str or not IDENTITY.fullmatch(value):
        raise ValueError("exact bounded identity required")
    return value


def sha256_text(value):
    if type(value) is not str or not SHA256.fullmatch(value):
        raise ValueError("exact 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 non-boolean integer required")
    return value


def finite_float(value, lower, upper, *, lower_inclusive=True):
    if type(value) is not float or not math.isfinite(value):
        raise ValueError("exact finite float required")
    lower_ok = value >= lower if lower_inclusive else value > lower
    if not lower_ok or value > upper:
        raise ValueError("bounded finite float required")
    return value


def exact_bool(value):
    if type(value) is not bool:
        raise ValueError("exact boolean required")
    return value


def bounded_sequence(value, lower, upper):
    if type(value) not in (tuple, list):
        raise ValueError("bounded tuple or list required")
    integer(len(value), lower, upper)
    return tuple(value)


def digest(value):
    try:
        encoded = json.dumps(
            value, sort_keys=True, separators=(",", ":"), allow_nan=False
        ).encode()
    except (TypeError, ValueError, OverflowError) as error:
        raise ValueError("canonical JSON value required") from error
    return hashlib.sha256(encoded).hexdigest()


def seal(record):
    if type(record.content_id) is not str:
        raise ValueError("content digest requires exact string")
    data = asdict(record)
    data.pop("content_id")
    return digest(data)


def validate_record(record, cls):
    if type(record) is not cls:
        raise ValueError("concrete frozen record required")
    try:
        rebuilt = cls(
            **{field.name: getattr(record, field.name) for field in fields(cls)}
        )
    except (AttributeError, TypeError, ValueError, OverflowError) as error:
        raise ValueError("malformed record") from error
    if record != rebuilt or record.content_id != rebuilt.content_id:
        raise ValueError("noncanonical or modified record")
    return rebuilt


def scope_tuple(value):
    result = bounded_sequence(value, 8, 8)
    for item in result:
        identity(item)
    if len(set(result)) != len(result):
        raise ValueError("duplicate scope identity")
    return result


def identity_tuple(value, lower, upper):
    result = bounded_sequence(value, lower, upper)
    for item in result:
        identity(item)
    return result


def rank_layout_tuple(value):
    result = bounded_sequence(value, 1, 1_000_000)
    for node in result:
        integer(node, 0, 99_999)
    return result


@dataclass(frozen=True)
class TrainingContract:
    scope: tuple = (
        "cluster-topology-v1",
        "model-revision-v1",
        "dataset-revision-v1",
        "optimizer-schema-v1",
        "precision-policy-v1",
        "collective-schema-v1",
        "checkpoint-schema-v1",
        "fixture-v1",
    )
    nodes: int = 2
    gpus_per_node: int = 8
    max_peak_memory_gib: float = 64.0
    minimum_step_throughput: float = 100.0
    max_collective_p95_ms: float = 50.0
    max_rank_step_skew: float = 0.10
    require_activation_checkpointing: bool = True
    require_checkpoint_restore: bool = True
    require_all_ranks_healthy: bool = True
    content_id: str = ""

    def __post_init__(self):
        object.__setattr__(self, "scope", scope_tuple(self.scope))
        integer(self.nodes, 1, 100_000)
        integer(self.gpus_per_node, 1, 100_000)
        if self.nodes * self.gpus_per_node > 1_000_000:
            raise ValueError("bounded training world size required")
        finite_float(self.max_peak_memory_gib, 0.0, 1_000_000.0, lower_inclusive=False)
        finite_float(
            self.minimum_step_throughput,
            0.0,
            1_000_000_000.0,
            lower_inclusive=False,
        )
        finite_float(self.max_collective_p95_ms, 0.0, 86_400_000.0, lower_inclusive=False)
        finite_float(self.max_rank_step_skew, 0.0, 1.0)
        for value in (
            self.require_activation_checkpointing,
            self.require_checkpoint_restore,
            self.require_all_ranks_healthy,
        ):
            if exact_bool(value) is not True:
                raise ValueError("required training control cannot be disabled")
        expected = seal(self)
        if self.content_id and self.content_id != expected:
            raise ValueError("training contract digest mismatch")
        object.__setattr__(self, "content_id", expected)


@dataclass(frozen=True)
class TrainingTopology:
    scope: tuple
    topology_id: str
    strategy: str
    data_parallel: int
    tensor_parallel: int
    pipeline_parallel: int
    parameter_shards: int
    gradient_shards: int
    optimizer_shards: int
    activation_checkpointing: bool
    collective_order: tuple
    content_id: str = ""

    def __post_init__(self):
        object.__setattr__(self, "scope", scope_tuple(self.scope))
        identity(self.topology_id)
        if type(self.strategy) is not str or self.strategy != "fully-sharded-data-parallel":
            raise ValueError("this audit requires fully-sharded-data-parallel")
        for value in (
            self.data_parallel,
            self.tensor_parallel,
            self.pipeline_parallel,
            self.parameter_shards,
            self.gradient_shards,
            self.optimizer_shards,
        ):
            integer(value, 1, 1_000_000)
        for shards in (
            self.parameter_shards,
            self.gradient_shards,
            self.optimizer_shards,
        ):
            if self.data_parallel % shards != 0:
                raise ValueError("state shard count must divide data-parallel degree")
        exact_bool(self.activation_checkpointing)
        order = identity_tuple(self.collective_order, 1, 16)
        object.__setattr__(self, "collective_order", order)
        expected = seal(self)
        if self.content_id and self.content_id != expected:
            raise ValueError("training topology digest mismatch")
        object.__setattr__(self, "content_id", expected)

    @property
    def world_size(self):
        return self.data_parallel * self.tensor_parallel * self.pipeline_parallel

    @property
    def mesh(self):
        return (
            f"dp{self.data_parallel}-tp{self.tensor_parallel}-"
            f"pp{self.pipeline_parallel}"
        )


@dataclass(frozen=True)
class TrainingEvidence:
    scope: tuple
    contract_content_id: str
    topology: TrainingTopology
    run_id: str
    observed_rank_count: int
    rank_to_node: tuple
    peak_memory_gib: float
    step_throughput: float
    collective_p95_ms: float
    rank_step_skew: float
    all_ranks_healthy: bool
    checkpoint_restore_verified: bool
    telemetry_complete: bool
    content_id: str = ""

    def __post_init__(self):
        object.__setattr__(self, "scope", scope_tuple(self.scope))
        sha256_text(self.contract_content_id)
        topology = validate_record(self.topology, TrainingTopology)
        if topology.scope != self.scope:
            raise ValueError("topology outside training scope")
        identity(self.run_id)
        integer(self.observed_rank_count, 1, 1_000_000)
        layout = rank_layout_tuple(self.rank_to_node)
        object.__setattr__(self, "rank_to_node", layout)
        finite_float(self.peak_memory_gib, 0.0, 1_000_000.0, lower_inclusive=False)
        finite_float(
            self.step_throughput, 0.0, 1_000_000_000.0, lower_inclusive=False
        )
        finite_float(self.collective_p95_ms, 0.0, 86_400_000.0, lower_inclusive=False)
        finite_float(self.rank_step_skew, 0.0, 1.0)
        exact_bool(self.all_ranks_healthy)
        exact_bool(self.checkpoint_restore_verified)
        exact_bool(self.telemetry_complete)
        object.__setattr__(self, "topology", topology)
        expected = seal(self)
        if self.content_id and self.content_id != expected:
            raise ValueError("training evidence digest mismatch")
        object.__setattr__(self, "content_id", expected)


@dataclass(frozen=True)
class TrainingAudit:
    decision: str
    violations: tuple
    topology: str
    world_size: int
    parameter_shards: int
    gradient_shards: int
    optimizer_shards: int
    peak_memory_gib: float
    step_throughput: float
    collective_p95_ms: float
    rank_step_skew: float
    evidence_id: str
    claim: str = "LOCAL_TRAINING_AUDIT_NOT_DISTRIBUTED_RUN"


def audit_training_topology(
    contract: TrainingContract, evidence: TrainingEvidence
) -> TrainingAudit:
    """Audit one supplied topology snapshot without launching distributed work."""
    contract = validate_record(contract, TrainingContract)
    evidence = validate_record(evidence, TrainingEvidence)
    if evidence.scope != contract.scope or evidence.contract_content_id != contract.content_id:
        raise ValueError("evidence belongs to another training contract")
    topology = evidence.topology
    world_size = contract.nodes * contract.gpus_per_node
    if topology.world_size != world_size:
        raise ValueError("training mesh does not cover the cluster world size")
    if evidence.observed_rank_count != world_size or len(evidence.rank_to_node) != world_size:
        raise ValueError("rank evidence does not cover the world size")
    rank_counts = Counter(evidence.rank_to_node)
    expected_rank_counts = {
        node: contract.gpus_per_node for node in range(contract.nodes)
    }
    if rank_counts != expected_rank_counts:
        raise ValueError("rank layout does not match GPUs per node")
    if (
        topology.parameter_shards,
        topology.gradient_shards,
        topology.optimizer_shards,
    ) != (topology.data_parallel,) * 3:
        raise ValueError("full sharding must partition parameters, gradients, and optimizer state")
    if topology.collective_order != (
        "parameter-all-gather",
        "gradient-reduce-scatter",
    ):
        raise ValueError("full-shard collective contract mismatch")

    readiness = []
    if contract.require_activation_checkpointing and not topology.activation_checkpointing:
        readiness.append("activation-checkpointing")
    if contract.require_checkpoint_restore and not evidence.checkpoint_restore_verified:
        readiness.append("checkpoint-restore")
    if contract.require_all_ranks_healthy and not evidence.all_ranks_healthy:
        readiness.append("rank-health")
    if not evidence.telemetry_complete:
        readiness.append("telemetry-incomplete")

    regressions = []
    if evidence.peak_memory_gib > contract.max_peak_memory_gib:
        regressions.append("peak-memory")
    if evidence.step_throughput < contract.minimum_step_throughput:
        regressions.append("step-throughput")
    if evidence.collective_p95_ms > contract.max_collective_p95_ms:
        regressions.append("collective-p95")
    if evidence.rank_step_skew > contract.max_rank_step_skew:
        regressions.append("rank-step-skew")

    violations = tuple(readiness + regressions)
    if regressions:
        decision = "BLOCK_TOPOLOGY"
    elif readiness:
        decision = "HOLD_INCOMPLETE"
    else:
        decision = "PASS_TOPOLOGY_AUDIT"
    evidence_id = digest(
        {
            "contract": contract.content_id,
            "evidence": evidence.content_id,
            "decision": decision,
            "violations": violations,
        }
    )
    return TrainingAudit(
        decision,
        violations,
        topology.mesh,
        world_size,
        topology.parameter_shards,
        topology.gradient_shards,
        topology.optimizer_shards,
        evidence.peak_memory_gib,
        evidence.step_throughput,
        evidence.collective_p95_ms,
        evidence.rank_step_skew,
        evidence_id,
    )


def illustrative_fixture():
    contract = TrainingContract()
    topology = TrainingTopology(
        contract.scope,
        "two-node-full-shard-v1",
        "fully-sharded-data-parallel",
        8,
        2,
        1,
        8,
        8,
        8,
        True,
        ("parameter-all-gather", "gradient-reduce-scatter"),
    )
    evidence = TrainingEvidence(
        contract.scope,
        contract.content_id,
        topology,
        "invented-training-run-v1",
        16,
        tuple([0] * 8 + [1] * 8),
        58.0,
        120.0,
        35.0,
        0.06,
        True,
        True,
        True,
    )
    return contract, evidence


def main():
    contract, evidence = illustrative_fixture()
    report = audit_training_topology(contract, evidence)
    print("example=illustrative_only")
    print(f"decision={report.decision}")
    print(f"topology={report.topology};world_size={report.world_size}")
    print(
        f"state_shards=parameters:{report.parameter_shards},"
        f"gradients:{report.gradient_shards},optimizer:{report.optimizer_shards}"
    )
    print(
        f"peak_memory_gib={report.peak_memory_gib:.3f};"
        f"step_throughput={report.step_throughput:.3f}"
    )
    print(
        f"collective_p95_ms={report.collective_p95_ms:.3f};"
        f"rank_step_skew={report.rank_step_skew:.3f}"
    )
    print(f"claim={report.claim}")


if __name__ == "__main__":
    main()