"""Diagnostic schedulability proof and claim-time authoritative recheck.

The proof is evidence, never a bearer credential.

A claim proceeds in three stages:
1. read-only preflight/discovery;
2. claim-time read-only recheck;
3. final *atomic* authority operation that revalidates all mutable claim
   conditions at the linearization point and returns lease/fencing receipt.

This module does not create a second scheduler or a second TaskPlane authority.
The host adapter must bind ``claim_atomically`` to the existing authoritative
claim transaction / repository operation.
"""
from __future__ import annotations

from dataclasses import dataclass, asdict
from typing import Any, Mapping, Protocol, Sequence
import hashlib
import time
import uuid

from .models import canonical_json


CLAIMABLE_TASK_STATES = frozenset({"CREATED", "QUEUED", "FAILED_RECOVERABLE"})


class SchedulabilityError(RuntimeError):
    pass


class SchedulabilityChanged(SchedulabilityError):
    pass


def _canonical(value: object) -> bytes:
    # Reuse the existing TaskPlane canonical JSON contract. Do not introduce
    # another sort/separator encoder for schedulability evidence.
    return canonical_json(value).encode("utf-8")


def _sha(value: object) -> str:
    return hashlib.sha256(_canonical(value)).hexdigest()


def _text(value: object, name: str, *, max_len: int = 256) -> str:
    if not isinstance(value, str):
        raise SchedulabilityError(f"{name} must be a string")
    value = value.strip()
    if not value or len(value) > max_len or any(ord(ch) < 32 for ch in value):
        raise SchedulabilityError(f"invalid {name}")
    return value


def _int_authority(
    value: object,
    *,
    name: str,
    default: int | None = None,
    minimum: int = 0,
) -> tuple[int, "Blocker | None"]:
    if value is None:
        if default is None:
            return 0, Blocker("INVALID_AUTHORITY_VALUE", "worker", f"{name} missing")
        return default, None
    if isinstance(value, bool) or not isinstance(value, int):
        return 0, Blocker("INVALID_AUTHORITY_VALUE", "worker", f"{name} must be integer")
    if value < minimum:
        return value, Blocker("INVALID_AUTHORITY_VALUE", "worker", f"{name} below {minimum}")
    return value, None


@dataclass(frozen=True)
class Blocker:
    code: str
    plane: str
    detail: str

    def to_dict(self) -> dict[str, str]:
        return asdict(self)


@dataclass(frozen=True)
class ClaimConditions:
    """Mutable authority conditions that MUST be checked atomically with claim."""

    task_revision: int
    task_generation: int
    task_status: str
    worker_generation: int
    worker_identity: str
    authorization_ref: str
    authorization_revision: str
    scope_digest: str
    capacity_limit: int

    def to_dict(self) -> dict[str, object]:
        return asdict(self)


@dataclass(frozen=True)
class SchedulabilityObservation:
    task_id: str
    task_revision: int
    task_generation: int
    task_status: str
    worker_id: str
    worker_generation: int
    worker_identity: str
    heartbeat_observed_at: float
    heartbeat_fresh: bool
    required_capabilities: tuple[str, ...]
    worker_capabilities: tuple[str, ...]
    scope_digest: str
    authorization_ref: str
    authorization_revision: str
    authorization_allowed: bool
    capacity_limit: int
    active_claims: int
    dependencies_satisfied: bool
    freeze_blocked: bool
    conflicting_lease: bool

    @property
    def capacity_remaining(self) -> int:
        return max(0, self.capacity_limit - self.active_claims)

    def stable_binding(self) -> dict[str, object]:
        """Fields whose change invalidates an earlier proof."""
        return {
            "task_id": self.task_id,
            "task_revision": self.task_revision,
            "task_generation": self.task_generation,
            "task_status": self.task_status,
            "worker_id": self.worker_id,
            "worker_generation": self.worker_generation,
            "worker_identity": self.worker_identity,
            "required_capabilities": list(self.required_capabilities),
            "worker_capabilities": list(self.worker_capabilities),
            "scope_digest": self.scope_digest,
            "authorization_ref": self.authorization_ref,
            "authorization_revision": self.authorization_revision,
            "authorization_allowed": self.authorization_allowed,
            "capacity_limit": self.capacity_limit,
            "active_claims": self.active_claims,
            "dependencies_satisfied": self.dependencies_satisfied,
            "freeze_blocked": self.freeze_blocked,
            "conflicting_lease": self.conflicting_lease,
            "heartbeat_fresh": self.heartbeat_fresh,
        }

    def claim_conditions(self) -> ClaimConditions:
        return ClaimConditions(
            task_revision=self.task_revision,
            task_generation=self.task_generation,
            task_status=self.task_status,
            worker_generation=self.worker_generation,
            worker_identity=self.worker_identity,
            authorization_ref=self.authorization_ref,
            authorization_revision=self.authorization_revision,
            scope_digest=self.scope_digest,
            capacity_limit=self.capacity_limit,
        )


@dataclass(frozen=True)
class SchedulabilityProof:
    schema: str
    proof_id: str
    task_id: str
    worker_id: str
    actor_identity: str
    created_at: float
    expires_at: float
    observation: SchedulabilityObservation
    blockers: tuple[Blocker, ...]
    observation_digest: str
    proof_digest: str

    @classmethod
    def create(
        cls,
        *,
        task_id: str,
        worker_id: str,
        actor_identity: str,
        observation: SchedulabilityObservation,
        blockers: Sequence[Blocker],
        now: float,
        ttl_seconds: float,
        proof_id: str | None = None,
    ) -> "SchedulabilityProof":
        if ttl_seconds <= 0 or ttl_seconds > 300:
            raise ValueError("schedulability proof ttl must be in (0,300]")
        payload = {
            "schema": "aios.schedulability-proof.v2",
            "proof_id": proof_id or f"sched-{uuid.uuid4()}",
            "task_id": _text(task_id, "task_id"),
            "worker_id": _text(worker_id, "worker_id"),
            "actor_identity": _text(actor_identity, "actor_identity", max_len=512),
            "created_at": float(now),
            "expires_at": float(now + ttl_seconds),
            "observation": observation.stable_binding(),
            "blockers": [b.to_dict() for b in blockers],
        }
        obs_digest = _sha(payload["observation"])
        proof_digest = _sha({**payload, "observation_digest": obs_digest})
        return cls(
            schema=payload["schema"],
            proof_id=payload["proof_id"],
            task_id=payload["task_id"],
            worker_id=payload["worker_id"],
            actor_identity=payload["actor_identity"],
            created_at=payload["created_at"],
            expires_at=payload["expires_at"],
            observation=observation,
            blockers=tuple(blockers),
            observation_digest=obs_digest,
            proof_digest=proof_digest,
        )

    @property
    def claimable(self) -> bool:
        return not self.blockers


@dataclass(frozen=True)
class DiscoveryDiagnostics:
    candidate_count: int
    blocker_codes: tuple[str, ...]


class SchedulabilityAuthority(Protocol):
    """Read authorities and perform the final existing claim transaction."""

    def read_task(self, task_id: str) -> Mapping[str, Any]: ...
    def list_candidate_tasks(self) -> Sequence[Mapping[str, Any]]: ...
    def read_worker(self, worker_id: str) -> Mapping[str, Any]: ...
    def read_authorization(
        self, *, task: Mapping[str, Any], worker: Mapping[str, Any], actor_identity: str
    ) -> Mapping[str, Any]: ...
    def active_claim_count(self, worker_id: str) -> int: ...
    def dependencies_satisfied(self, task: Mapping[str, Any]) -> bool: ...
    def freeze_blocked(self, task: Mapping[str, Any]) -> bool: ...
    def conflicting_lease(self, task_id: str) -> bool: ...
    def exact_scope_digest(
        self, *, task: Mapping[str, Any], worker: Mapping[str, Any], actor_identity: str
    ) -> tuple[bool, str]: ...
    def claim_atomically(
        self,
        *,
        task_id: str,
        worker_id: str,
        actor_identity: str,
        expected: ClaimConditions,
        lease_seconds: int,
        schedulability_proof_digest: str,
    ) -> Mapping[str, Any]:
        """Claim at the existing TaskPlane linearization point.

        The implementation MUST, in the same authoritative transaction / fenced
        operation that writes the claim:
        - verify task revision/generation/status is still claimable;
        - verify worker registration/identity/health/capacity;
        - verify authorization ref/revision remains allowed;
        - verify exact scope;
        - verify dependencies/freeze/conflicting lease;
        - create claim + lease/fencing receipt.

        This protocol is an adapter contract; it is not evidence that a specific
        production adapter already implements the transaction.
        """


class ProofRepository(Protocol):
    def put(self, proof: SchedulabilityProof) -> None: ...
    def get(self, proof_id: str) -> SchedulabilityProof | None: ...


class MemoryProofRepository:
    def __init__(self) -> None:
        self._values: dict[str, SchedulabilityProof] = {}

    def put(self, proof: SchedulabilityProof) -> None:
        existing = self._values.get(proof.proof_id)
        if existing is not None and existing.proof_digest != proof.proof_digest:
            raise SchedulabilityError("proof id replay conflict")
        self._values[proof.proof_id] = proof

    def get(self, proof_id: str) -> SchedulabilityProof | None:
        return self._values.get(proof_id)


class SchedulabilityService:
    def __init__(
        self,
        authority: SchedulabilityAuthority,
        proofs: ProofRepository,
        *,
        clock=time.time,
        proof_ttl_seconds: float = 30.0,
        heartbeat_grace_factor: float = 3.0,
    ) -> None:
        self.authority = authority
        self.proofs = proofs
        self.clock = clock
        self.proof_ttl_seconds = proof_ttl_seconds
        self.heartbeat_grace_factor = heartbeat_grace_factor

    @staticmethod
    def _caps(value: object) -> tuple[str, ...]:
        if not isinstance(value, (list, tuple)):
            return ()
        return tuple(sorted({str(v) for v in value if isinstance(v, str) and v}))

    def observe(
        self, task_id: str, worker_id: str, actor_identity: str
    ) -> tuple[SchedulabilityObservation, tuple[Blocker, ...]]:
        task = self.authority.read_task(task_id)
        worker = self.authority.read_worker(worker_id)
        auth = self.authority.read_authorization(
            task=task, worker=worker, actor_identity=actor_identity
        )

        required = self._caps(task.get("required_capabilities") or ())
        offered = self._caps(worker.get("capabilities") or ())
        now = float(self.clock())
        last = float(worker.get("last_heartbeat") or worker.get("last_seen_at_epoch") or 0.0)
        interval = float(worker.get("heartbeat_interval_seconds") or 15.0)
        fresh = (
            str(worker.get("status") or "").lower() == "active"
            and not bool(worker.get("disabled"))
            and last <= now
            and now - last <= max(interval * self.heartbeat_grace_factor, 1.0)
        )
        scope_ok, scope_digest = self.authority.exact_scope_digest(
            task=task, worker=worker, actor_identity=actor_identity
        )
        deps = bool(self.authority.dependencies_satisfied(task))
        freeze = bool(self.authority.freeze_blocked(task))
        lease = bool(self.authority.conflicting_lease(task_id))
        active_raw = self.authority.active_claim_count(worker_id)
        active, active_error = _int_authority(active_raw, name="active_claim_count", minimum=0)
        capacity, capacity_error = _int_authority(
            worker.get("max_concurrency"),
            name="max_concurrency",
            default=1,
            minimum=0,
        )
        auth_allowed = bool(auth.get("allowed"))
        task_status = str(task.get("status") or "")
        blockers: list[Blocker] = []

        if task_status not in CLAIMABLE_TASK_STATES:
            blockers.append(
                Blocker(
                    "TASK_STATE_NOT_CLAIMABLE",
                    "task",
                    f"task status {task_status or '<missing>'} is not claimable",
                )
            )
        if active_error:
            blockers.append(active_error)
        if capacity_error:
            blockers.append(capacity_error)
        if not fresh:
            blockers.append(Blocker("STALE_OR_INACTIVE_WORKER", "worker", "worker heartbeat is not fresh"))
        missing = sorted(set(required) - set(offered))
        if missing:
            blockers.append(Blocker("MISSING_CAPABILITY", "worker", ",".join(missing)))
        if not scope_ok:
            blockers.append(Blocker("EXACT_SCOPE_MISMATCH", "authorization", "worker/task signed scope mismatch"))
        if not auth_allowed:
            blockers.append(Blocker("AUTHORIZATION_BLOCKED", "authorization", "current authorization does not allow claim"))
        if capacity >= 0 and active >= capacity:
            blockers.append(Blocker("CAPACITY_EXHAUSTED", "worker", f"{active}/{capacity} active"))
        if not deps:
            blockers.append(Blocker("DEPENDENCY_BLOCKED", "task", "task dependencies are not satisfied"))
        if freeze:
            blockers.append(Blocker("FREEZE_BLOCKED", "task", "task/environment freeze is active"))
        if lease:
            blockers.append(Blocker("CONFLICTING_LEASE", "task", "a live conflicting lease exists"))

        obs = SchedulabilityObservation(
            task_id=_text(task.get("task_id"), "task_id"),
            task_revision=int(task.get("revision") or 0),
            task_generation=int(task.get("generation") or task.get("metadata", {}).get("generation") or 0),
            task_status=task_status,
            worker_id=_text(worker.get("worker_id"), "worker_id"),
            worker_generation=int(worker.get("registration_revision") or worker.get("generation") or 0),
            worker_identity=_text(worker.get("identity"), "worker_identity", max_len=512),
            heartbeat_observed_at=last,
            heartbeat_fresh=fresh,
            required_capabilities=required,
            worker_capabilities=offered,
            scope_digest=_text(scope_digest, "scope_digest", max_len=256),
            authorization_ref=_text(auth.get("authorization_ref"), "authorization_ref", max_len=512),
            authorization_revision=_text(auth.get("revision"), "authorization_revision", max_len=256),
            authorization_allowed=auth_allowed,
            capacity_limit=capacity,
            active_claims=active,
            dependencies_satisfied=deps,
            freeze_blocked=freeze,
            conflicting_lease=lease,
        )
        return obs, tuple(blockers)

    def preflight(
        self,
        *,
        task_id: str,
        worker_id: str,
        actor_identity: str,
        persist: bool = False,
    ) -> SchedulabilityProof:
        now = float(self.clock())
        obs, blockers = self.observe(task_id, worker_id, actor_identity)
        proof = SchedulabilityProof.create(
            task_id=task_id,
            worker_id=worker_id,
            actor_identity=actor_identity,
            observation=obs,
            blockers=blockers,
            now=now,
            ttl_seconds=self.proof_ttl_seconds,
        )
        if persist:
            self.proofs.put(proof)
        return proof

    def discover_with_diagnostics(
        self, *, worker_id: str, actor_identity: str
    ) -> tuple[SchedulabilityProof | None, DiscoveryDiagnostics]:
        candidates = sorted(
            self.authority.list_candidate_tasks(),
            key=lambda task: (
                -int(task.get("priority") or 0),
                str(task.get("created_at") or ""),
                str(task.get("task_id") or ""),
            ),
        )
        codes: set[str] = set()
        observed = 0
        for task in candidates:
            task_id = str(task.get("task_id") or "")
            if not task_id:
                continue
            observed += 1
            proof = self.preflight(
                task_id=task_id,
                worker_id=worker_id,
                actor_identity=actor_identity,
                persist=False,
            )
            if proof.claimable:
                self.proofs.put(proof)
                return proof, DiscoveryDiagnostics(observed, tuple(sorted(codes)))
            codes.update(b.code for b in proof.blockers)
        return None, DiscoveryDiagnostics(observed, tuple(sorted(codes)))

    def discover(
        self, *, worker_id: str, actor_identity: str
    ) -> SchedulabilityProof | None:
        proof, _ = self.discover_with_diagnostics(
            worker_id=worker_id, actor_identity=actor_identity
        )
        return proof

    @staticmethod
    def _validate_claim_receipt(result: Mapping[str, Any]) -> None:
        required = ("claim_id", "lease_id", "fencing_generation", "receipt_revision")
        missing = [name for name in required if result.get(name) in (None, "")]
        if missing:
            raise SchedulabilityChanged(
                "claim authority did not return required lease/fence receipt: "
                + ",".join(missing)
            )

    def claim_with_recheck(
        self,
        *,
        proof_id: str,
        worker_id: str,
        actor_identity: str,
        lease_seconds: int = 30,
    ) -> Mapping[str, Any]:
        proof = self.proofs.get(proof_id)
        if proof is None:
            raise SchedulabilityChanged("proof not found")
        now = float(self.clock())
        if proof.expires_at <= now:
            raise SchedulabilityChanged("proof expired")
        if proof.worker_id != worker_id or proof.actor_identity != actor_identity:
            raise SchedulabilityChanged("proof identity mismatch")
        current, blockers = self.observe(proof.task_id, worker_id, actor_identity)
        if blockers:
            raise SchedulabilityChanged(
                "claim-time blockers: " + ",".join(sorted(b.code for b in blockers))
            )
        if current.stable_binding() != proof.observation.stable_binding():
            raise SchedulabilityChanged("claim-time authority changed since proof")

        result = self.authority.claim_atomically(
            task_id=proof.task_id,
            worker_id=worker_id,
            actor_identity=actor_identity,
            expected=current.claim_conditions(),
            lease_seconds=lease_seconds,
            schedulability_proof_digest=proof.proof_digest,
        )
        self._validate_claim_receipt(result)
        return result