from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass
from typing import Any

from pydantic import BaseModel, ConfigDict, Field, model_validator

from .admission import AdmissionDecision, AdmissionGate, SignedAdmissionToken
from .genome import ServiceGenome
from .ledger import ReplayTrace, TraceDelta
from .policy import SafetyAttestation


class ReplicaWeaveIntent(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    source: str
    target: str
    intent_kind: str
    budget_hint: dict[str, Any] = Field(default_factory=dict)
    safety_labels: list[str] = Field(default_factory=list)

    @model_validator(mode="after")
    def _validate_identity(self) -> "ReplicaWeaveIntent":
        if not self.source.strip() or not self.target.strip():
            raise ValueError("source and target are required")
        return self


class RolloutStage(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    name: str
    cohort_size: int
    percentage: int


class ReplicationPlan(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    genome_hash: str
    approved: bool
    stages: list[RolloutStage] = Field(default_factory=list)
    trace: ReplayTrace
    rationale: list[str] = Field(default_factory=list)
    admission_decision: AdmissionDecision | None = None

    def digest(self) -> str:
        payload = self.model_dump(mode="json", exclude_none=True)
        return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True).encode("utf-8")).hexdigest()


def _rollout_stages(max_replicas: int, score: float) -> list[RolloutStage]:
    canary = min(1, max_replicas)
    cohort = max(1, min(max_replicas, max(2, max_replicas // 2)))
    stages = [
        RolloutStage(name="canary", cohort_size=canary, percentage=5),
        RolloutStage(name="cohort", cohort_size=cohort, percentage=25),
        RolloutStage(name="regional", cohort_size=max_replicas, percentage=100),
    ]
    return stages[: 2 if score < 85 else 3]


class ReplicatorPlanner:
    def plan(
        self,
        genome: ServiceGenome,
        attestation: SafetyAttestation,
        *,
        intent: ReplicaWeaveIntent | None = None,
        admission_token: SignedAdmissionToken | None = None,
    ) -> ReplicationPlan:
        genome_hash = genome.content_hash()
        trace_seed = int(genome_hash[:16], 16)
        rationale = list(attestation.reasons)
        admission_decision = None

        if admission_token is not None:
            admission_decision = AdmissionGate().authorize(genome, admission_token, intent=intent)
            if not admission_decision.approved:
                rationale.extend(admission_decision.reasons)
                trace = ReplayTrace(
                    seed=trace_seed,
                    deltas=[
                        TraceDelta(
                            stage="blocked",
                            operation="deny",
                            payload={
                                "policy": attestation.policy_name,
                                "admission_token": admission_decision.token_digest,
                                "reasons": list(admission_decision.reasons),
                            },
                        )
                    ],
                )
                return ReplicationPlan(
                    genome_hash=genome_hash,
                    approved=False,
                    stages=[],
                    trace=trace,
                    rationale=rationale,
                    admission_decision=admission_decision,
                )

        if not attestation.approved:
            trace = ReplayTrace(
                seed=trace_seed,
                deltas=[TraceDelta(stage="blocked", operation="deny", payload={"policy": attestation.policy_name})],
            )
            return ReplicationPlan(
                genome_hash=genome_hash,
                approved=False,
                stages=[],
                trace=trace,
                rationale=rationale,
                admission_decision=admission_decision,
            )

        stages = _rollout_stages(genome.resource_envelope.max_replicas, attestation.score)
        deltas = [
            TraceDelta(stage=stage.name, operation="rollout", payload={"cohort_size": stage.cohort_size, "percentage": stage.percentage})
            for stage in stages
        ]
        signatures = [genome_hash, attestation.digest()]
        if admission_decision is not None:
            signatures.append(admission_decision.digest())
        trace = ReplayTrace(seed=trace_seed, deltas=deltas, signatures=signatures)
        return ReplicationPlan(
            genome_hash=genome_hash,
            approved=True,
            stages=stages,
            trace=trace,
            rationale=rationale,
            admission_decision=admission_decision,
        )


class ReplayHarness:
    def verify(
        self,
        genome: ServiceGenome,
        attestation: SafetyAttestation,
        plan: ReplicationPlan,
        *,
        intent: ReplicaWeaveIntent | None = None,
        admission_token: SignedAdmissionToken | None = None,
    ) -> ReplayTrace:
        genome_hash = genome.content_hash()
        if plan.genome_hash != genome_hash:
            raise ValueError("plan genome hash does not match genome content")

        expected_seed = int(genome_hash[:16], 16)
        plan.trace.assert_valid(seed=expected_seed)

        if plan.approved != attestation.approved:
            raise ValueError("plan approval does not match attestation")

        expected_admission = None
        if admission_token is not None:
            expected_admission = AdmissionGate().authorize(genome, admission_token, intent=intent)
            if plan.admission_decision != expected_admission:
                raise ValueError("plan admission decision is not reproducible from the token and intent")
            if expected_admission.approved and expected_admission.digest() not in plan.trace.signatures:
                raise ValueError("trace signatures do not include the admission decision")

        if not plan.approved:
            expected_payload: dict[str, object] = {"policy": attestation.policy_name}
            if expected_admission is not None and not expected_admission.approved:
                expected_payload["admission_token"] = expected_admission.token_digest
                expected_payload["reasons"] = list(expected_admission.reasons)
            if plan.trace.deltas != [TraceDelta(stage="blocked", operation="deny", payload=expected_payload)]:
                raise ValueError("blocked trace does not match attestation")
            return plan.trace

        expected_stages = _rollout_stages(genome.resource_envelope.max_replicas, attestation.score)
        expected_deltas = [
            TraceDelta(stage=stage.name, operation="rollout", payload={"cohort_size": stage.cohort_size, "percentage": stage.percentage})
            for stage in expected_stages
        ]

        if plan.stages != expected_stages:
            raise ValueError("plan stages are not reproducible from the genome and attestation")
        if plan.trace.deltas != expected_deltas:
            raise ValueError("trace deltas are not reproducible from the rollout stages")
        expected_signatures = [genome_hash, attestation.digest()]
        if expected_admission is not None:
            expected_signatures.append(expected_admission.digest())
        if plan.trace.signatures != expected_signatures:
            raise ValueError("trace signatures do not match the expected audit chain")
        return plan.trace