"""Deterministic inference-parallelism planner over invented benchmark evidence."""

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, 7, 7)
    for item in result:
        identity(item)
    if len(set(result)) != len(result):
        raise ValueError("duplicate scope identity")
    return result


@dataclass(frozen=True)
class ParallelismContract:
    scope: tuple = (
        "cluster-topology-v1",
        "model-profile-v1",
        "benchmark-schema-v1",
        "memory-policy-v1",
        "latency-policy-v1",
        "throughput-policy-v1",
        "fixture-v1",
    )
    nodes: int = 2
    gpus_per_node: int = 8
    per_gpu_memory_gib: float = 80.0
    reserved_memory_gib: float = 8.0
    max_p95_ttft_ms: float = 500.0
    minimum_tokens_per_second: float = 900.0
    max_inter_node_link_utilization: float = 0.80
    require_tensor_parallel_within_node: bool = True
    require_complete_telemetry: 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 cluster world size required")
        finite_float(self.per_gpu_memory_gib, 0.0, 1_000_000.0, lower_inclusive=False)
        finite_float(self.reserved_memory_gib, 0.0, self.per_gpu_memory_gib)
        if self.reserved_memory_gib >= self.per_gpu_memory_gib:
            raise ValueError("reserved memory must leave usable capacity")
        finite_float(self.max_p95_ttft_ms, 0.0, 86_400_000.0, lower_inclusive=False)
        finite_float(
            self.minimum_tokens_per_second,
            0.0,
            1_000_000_000.0,
            lower_inclusive=False,
        )
        finite_float(self.max_inter_node_link_utilization, 0.0, 1.0)
        exact_bool(self.require_tensor_parallel_within_node)
        if exact_bool(self.require_complete_telemetry) is not True:
            raise ValueError("complete benchmark telemetry is required")
        expected = seal(self)
        if self.content_id and self.content_id != expected:
            raise ValueError("parallelism contract digest mismatch")
        object.__setattr__(self, "content_id", expected)


@dataclass(frozen=True)
class ModelProfile:
    scope: tuple
    model_id: str
    architecture: str
    parameter_billions: float
    active_parameter_billions: float
    expert_count: int
    max_context_tokens: int
    content_id: str = ""

    def __post_init__(self):
        object.__setattr__(self, "scope", scope_tuple(self.scope))
        identity(self.model_id)
        if type(self.architecture) is not str or self.architecture not in (
            "dense",
            "moe",
        ):
            raise ValueError("architecture must be dense or moe")
        finite_float(
            self.parameter_billions, 0.0, 100_000.0, lower_inclusive=False
        )
        finite_float(
            self.active_parameter_billions, 0.0, self.parameter_billions,
            lower_inclusive=False,
        )
        integer(self.expert_count, 1, 1_000_000)
        integer(self.max_context_tokens, 1, 100_000_000)
        if self.architecture == "dense" and (
            self.expert_count != 1
            or self.active_parameter_billions != self.parameter_billions
        ):
            raise ValueError("dense profile cannot declare sparse experts")
        if self.architecture == "moe" and self.expert_count < 2:
            raise ValueError("moe profile requires multiple experts")
        expected = seal(self)
        if self.content_id and self.content_id != expected:
            raise ValueError("model profile digest mismatch")
        object.__setattr__(self, "content_id", expected)


@dataclass(frozen=True)
class TopologyTrial:
    scope: tuple
    benchmark_id: str
    model_content_id: str
    tensor_parallel: int
    pipeline_parallel: int
    expert_parallel: int
    data_parallel: int
    tensor_group_within_node: bool
    peak_memory_gib: float
    p95_ttft_ms: float
    tokens_per_second: float
    inter_node_link_utilization: float
    telemetry_complete: bool
    content_id: str = ""

    def __post_init__(self):
        object.__setattr__(self, "scope", scope_tuple(self.scope))
        identity(self.benchmark_id)
        sha256_text(self.model_content_id)
        for value in (
            self.tensor_parallel,
            self.pipeline_parallel,
            self.expert_parallel,
            self.data_parallel,
        ):
            integer(value, 1, 1_000_000)
        exact_bool(self.tensor_group_within_node)
        finite_float(self.peak_memory_gib, 0.0, 1_000_000.0, lower_inclusive=False)
        finite_float(self.p95_ttft_ms, 0.0, 86_400_000.0, lower_inclusive=False)
        finite_float(
            self.tokens_per_second, 0.0, 1_000_000_000.0, lower_inclusive=False
        )
        finite_float(self.inter_node_link_utilization, 0.0, 1.0)
        exact_bool(self.telemetry_complete)
        expected = seal(self)
        if self.content_id and self.content_id != expected:
            raise ValueError("topology trial digest mismatch")
        object.__setattr__(self, "content_id", expected)

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

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


@dataclass(frozen=True)
class ParallelismEvidence:
    scope: tuple
    contract_content_id: str
    model: ModelProfile
    trials: tuple
    content_id: str = ""

    def __post_init__(self):
        object.__setattr__(self, "scope", scope_tuple(self.scope))
        sha256_text(self.contract_content_id)
        model = validate_record(self.model, ModelProfile)
        raw_trials = bounded_sequence(self.trials, 1, 1_000)
        trials = tuple(validate_record(item, TopologyTrial) for item in raw_trials)
        if model.scope != self.scope or any(item.scope != self.scope for item in trials):
            raise ValueError("parallelism evidence outside scope")
        if any(item.model_content_id != model.content_id for item in trials):
            raise ValueError("trial belongs to another model profile")
        if len({item.benchmark_id for item in trials}) != 1:
            raise ValueError("trials belong to different benchmark runs")
        shapes = {
            (
                item.tensor_parallel,
                item.pipeline_parallel,
                item.expert_parallel,
                item.data_parallel,
            )
            for item in trials
        }
        if len(shapes) != len(trials):
            raise ValueError("duplicate topology trial")
        object.__setattr__(self, "model", model)
        object.__setattr__(self, "trials", trials)
        expected = seal(self)
        if self.content_id and self.content_id != expected:
            raise ValueError("parallelism evidence digest mismatch")
        object.__setattr__(self, "content_id", expected)


@dataclass(frozen=True)
class ParallelismPlan:
    decision: str
    rejected: tuple
    eligible_count: int
    topology: str
    world_size: int
    nodes: int
    gpus_per_node: int
    peak_memory_gib: float
    p95_ttft_ms: float
    tokens_per_second: float
    inter_node_link_utilization: float
    evidence_id: str
    claim: str = "LOCAL_TOPOLOGY_AUDIT_NOT_PRODUCTION_BENCHMARK"


def plan_parallelism(
    contract: ParallelismContract, evidence: ParallelismEvidence
) -> ParallelismPlan:
    """Choose among supplied invented trials; never provision or benchmark hardware."""
    contract = validate_record(contract, ParallelismContract)
    evidence = validate_record(evidence, ParallelismEvidence)
    if evidence.scope != contract.scope or evidence.contract_content_id != contract.content_id:
        raise ValueError("evidence belongs to another parallelism contract")
    world_size = contract.nodes * contract.gpus_per_node
    candidates = []
    rejected = []
    for trial in evidence.trials:
        if trial.world_size != world_size:
            raise ValueError("topology degrees do not cover the cluster world size")
        if trial.tensor_group_within_node:
            if trial.tensor_parallel > contract.gpus_per_node:
                raise ValueError("within-node tensor group exceeds node capacity")
            if contract.gpus_per_node % trial.tensor_parallel != 0:
                raise ValueError("within-node tensor groups must pack evenly on each node")
        if evidence.model.architecture == "dense" and trial.expert_parallel != 1:
            raise ValueError("dense model cannot use expert parallelism")
        if evidence.model.expert_count % trial.expert_parallel != 0:
            raise ValueError("expert parallel degree must divide expert count")

        reasons = []
        if (
            contract.require_tensor_parallel_within_node
            and not trial.tensor_group_within_node
        ):
            reasons.append("tensor-group-placement")
        if not trial.telemetry_complete:
            reasons.append("telemetry-incomplete")
        if trial.peak_memory_gib > (
            contract.per_gpu_memory_gib - contract.reserved_memory_gib
        ):
            reasons.append("memory-headroom")
        if trial.p95_ttft_ms > contract.max_p95_ttft_ms:
            reasons.append("p95-ttft")
        if trial.tokens_per_second < contract.minimum_tokens_per_second:
            reasons.append("throughput")
        if (
            trial.inter_node_link_utilization
            > contract.max_inter_node_link_utilization
        ):
            reasons.append("inter-node-link")
        if reasons:
            rejected.extend(f"{trial.topology}:{reason}" for reason in reasons)
        else:
            candidates.append(trial)

    if candidates:
        chosen = min(
            candidates,
            key=lambda item: (
                -item.tokens_per_second,
                item.p95_ttft_ms,
                item.inter_node_link_utilization,
                item.topology,
            ),
        )
        decision = "SELECT_TOPOLOGY"
        topology = chosen.topology
        peak_memory = chosen.peak_memory_gib
        ttft = chosen.p95_ttft_ms
        throughput = chosen.tokens_per_second
        link = chosen.inter_node_link_utilization
    else:
        decision = "NO_ELIGIBLE_TOPOLOGY"
        topology = "none"
        peak_memory = 0.0
        ttft = 0.0
        throughput = 0.0
        link = 0.0
    evidence_id = digest(
        {
            "contract": contract.content_id,
            "evidence": evidence.content_id,
            "decision": decision,
            "topology": topology,
            "rejected": tuple(rejected),
        }
    )
    return ParallelismPlan(
        decision,
        tuple(rejected),
        len(candidates),
        topology,
        world_size,
        contract.nodes,
        contract.gpus_per_node,
        peak_memory,
        ttft,
        throughput,
        link,
        evidence_id,
    )


def illustrative_fixture():
    contract = ParallelismContract()
    model = ModelProfile(
        contract.scope,
        "invented-moe-64b",
        "moe",
        64.0,
        12.0,
        64,
        32_768,
    )
    trials = (
        TopologyTrial(
            contract.scope,
            "benchmark-run-v1",
            model.content_id,
            8,
            2,
            1,
            1,
            True,
            54.0,
            450.0,
            920.0,
            0.72,
            True,
        ),
        TopologyTrial(
            contract.scope,
            "benchmark-run-v1",
            model.content_id,
            4,
            2,
            2,
            1,
            True,
            50.0,
            420.0,
            960.0,
            0.60,
            True,
        ),
        TopologyTrial(
            contract.scope,
            "benchmark-run-v1",
            model.content_id,
            4,
            1,
            2,
            2,
            True,
            58.0,
            390.0,
            880.0,
            0.55,
            True,
        ),
    )
    evidence = ParallelismEvidence(
        contract.scope, contract.content_id, model, trials
    )
    return contract, evidence


def main():
    contract, evidence = illustrative_fixture()
    report = plan_parallelism(contract, evidence)
    print("example=illustrative_only")
    print(f"decision={report.decision}")
    print(f"topology={report.topology}")
    print(
        f"world_size={report.world_size};nodes={report.nodes};"
        f"gpus_per_node={report.gpus_per_node}"
    )
    print(
        f"peak_memory_gib={report.peak_memory_gib:.3f};"
        f"ttft_ms={report.p95_ttft_ms:.3f}"
    )
    print(
        f"tokens_per_second={report.tokens_per_second:.3f};"
        f"inter_node_link_utilization={report.inter_node_link_utilization:.3f}"
    )
    print(f"claim={report.claim}")


if __name__ == "__main__":
    main()