"""Versioned Shard lifecycle and structured status contract.

This module defines the semantic contract consumed by a future generated gRPC
binding. It deliberately contains no Meshnet routing, authentication policy,
billing, or llama.cpp types.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
import json
import threading
from typing import Any, Callable, ClassVar


SCHEMA_NAME = "meshnet.shard-runtime"
SCHEMA_VERSION = 1
SUPPORTED_VERSIONS = frozenset({SCHEMA_VERSION})
DEFAULT_MAX_INFLIGHT_FRAMES = 32
DEFAULT_MAX_INFLIGHT_BYTES = 8 * 1024 * 1024


class RpcName(str, Enum):
    CAPABILITY = "GetCapability"
    HEALTH = "CheckHealth"
    SESSION = "OpenSession"
    CANCEL = "CancelSession"
    RELEASE = "ReleaseSession"
    METRICS = "GetMetrics"


class StatusCode(str, Enum):
    OK = "OK"
    INVALID_ARGUMENT = "INVALID_ARGUMENT"
    UNSUPPORTED_VERSION = "UNSUPPORTED_VERSION"
    FAILED_PRECONDITION = "FAILED_PRECONDITION"
    MALFORMED_LIFECYCLE = "MALFORMED_LIFECYCLE"
    NOT_FOUND = "NOT_FOUND"
    ALREADY_EXISTS = "ALREADY_EXISTS"
    CANCELLED = "CANCELLED"
    DEADLINE_EXCEEDED = "DEADLINE_EXCEEDED"
    RESOURCE_EXHAUSTED = "RESOURCE_EXHAUSTED"
    UNAUTHENTICATED = "UNAUTHENTICATED"
    PERMISSION_DENIED = "PERMISSION_DENIED"
    DATA_LOSS = "DATA_LOSS"
    UNAVAILABLE = "UNAVAILABLE"
    INTERNAL = "INTERNAL"


class LifecycleState(str, Enum):
    OPEN = "OPEN"
    PREFILLING = "PREFILLING"
    DECODING = "DECODING"
    CANCELLING = "CANCELLING"
    CANCELLED = "CANCELLED"
    RELEASING = "RELEASING"
    RELEASED = "RELEASED"
    FAILED = "FAILED"


class CacheExpectation(str, Enum):
    NONE = "NONE"
    OPTIONAL = "OPTIONAL"
    REQUIRED = "REQUIRED"


class CacheResult(str, Enum):
    NOT_REQUESTED = "NOT_REQUESTED"
    HIT = "HIT"
    MISS = "MISS"
    INVALIDATED = "INVALIDATED"
    STORED = "STORED"


class SessionPhase(str, Enum):
    PREFILL = "PREFILL"
    DECODE = "DECODE"


class LifecycleContractError(ValueError):
    """A protocol violation represented by a structured status."""

    def __init__(self, status: "StructuredStatus") -> None:
        self.status = status
        super().__init__(status.message)


@dataclass(frozen=True)
class StructuredStatus:
    code: StatusCode
    message: str
    retryable: bool = False
    details: dict[str, str] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "code": self.code.value,
            "message": self.message,
            "retryable": self.retryable,
            "details": dict(sorted(self.details.items())),
        }

    @classmethod
    def from_dict(cls, value: dict[str, Any]) -> "StructuredStatus":
        return cls(
            code=StatusCode(str(value["code"])),
            message=str(value["message"]),
            retryable=bool(value.get("retryable", False)),
            details={str(k): str(v) for k, v in value.get("details", {}).items()},
        )


@dataclass(frozen=True)
class DeadlinePolicy:
    capability_seconds: float = 5.0
    health_seconds: float = 2.0
    session_open_seconds: float = 10.0
    session_idle_seconds: float = 30.0
    cancel_seconds: float = 2.0
    release_seconds: float = 5.0
    metrics_seconds: float = 5.0

    def __post_init__(self) -> None:
        if any(value <= 0 for value in self.__dict__.values()):
            raise ValueError("all RPC deadlines must be positive")


@dataclass(frozen=True)
class TlsAuthHooks:
    """Transport hooks only; Meshnet remains the identity and billing authority."""

    tls_required: bool = True
    server_name: str = "shard.meshnet"
    client_certificate_hook: str | None = None
    peer_identity_hook: str | None = None
    auth_metadata_hook: str | None = None

    def validate(self) -> None:
        if self.tls_required and not self.server_name:
            raise ValueError("TLS server_name is required when TLS is enabled")


@dataclass(frozen=True)
class FlowControlLimits:
    max_inflight_frames: int = DEFAULT_MAX_INFLIGHT_FRAMES
    max_inflight_bytes: int = DEFAULT_MAX_INFLIGHT_BYTES
    max_frame_bytes: int = 1024 * 1024

    def __post_init__(self) -> None:
        if min(self.max_inflight_frames, self.max_inflight_bytes, self.max_frame_bytes) <= 0:
            raise ValueError("flow-control limits must be positive")


class FlowControl:
    """Bounded sender window; callers block or fail instead of growing unbounded."""

    def __init__(self, limits: FlowControlLimits = FlowControlLimits()) -> None:
        self.limits = limits
        self._condition = threading.Condition()
        self._frames = 0
        self._bytes = 0

    @property
    def outstanding(self) -> tuple[int, int]:
        with self._condition:
            return self._frames, self._bytes

    def acquire(self, size: int, *, wait: bool = False, cancelled: Callable[[], bool] | None = None) -> None:
        if size < 0 or size > self.limits.max_frame_bytes:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.RESOURCE_EXHAUSTED, "frame exceeds bounded flow-control window",
                details={"max_frame_bytes": str(self.limits.max_frame_bytes)},
            ))
        with self._condition:
            while (
                self._frames >= self.limits.max_inflight_frames
                or self._bytes + size > self.limits.max_inflight_bytes
            ):
                if cancelled and cancelled():
                    raise LifecycleContractError(StructuredStatus(StatusCode.CANCELLED, "flow-control wait cancelled"))
                if not wait:
                    raise LifecycleContractError(StructuredStatus(
                        StatusCode.RESOURCE_EXHAUSTED, "flow-control window is full", retryable=True,
                    ))
                self._condition.wait(timeout=0.05)
            self._frames += 1
            self._bytes += size

    def release(self, size: int) -> None:
        with self._condition:
            self._frames = max(0, self._frames - 1)
            self._bytes = max(0, self._bytes - max(0, size))
            self._condition.notify_all()


class CancellationToken:
    def __init__(self) -> None:
        self._event = threading.Event()

    @property
    def cancelled(self) -> bool:
        return self._event.is_set()

    def cancel(self) -> None:
        self._event.set()


@dataclass(frozen=True)
class CapabilityRequest:
    schema_version: int = SCHEMA_VERSION


@dataclass(frozen=True)
class CapabilityResponse:
    status: StructuredStatus
    schema_version: int
    supported_versions: tuple[int, ...]
    artifact_fingerprint: str = ""
    runtime_fingerprint: str = ""
    shard_start: int = 0
    shard_end: int = 0
    effective_start: int = 0
    max_sessions: int = 0
    max_frame_bytes: int = DEFAULT_MAX_INFLIGHT_BYTES

    def validate(self) -> None:
        if self.status.code is StatusCode.OK and self.schema_version not in SUPPORTED_VERSIONS:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.UNSUPPORTED_VERSION, "worker selected an unsupported schema version",
            ))
        if self.shard_end < self.shard_start or self.effective_start < self.shard_start:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.INVALID_ARGUMENT, "invalid authoritative shard range",
            ))


@dataclass(frozen=True)
class HealthRequest:
    schema_version: int = SCHEMA_VERSION
    include_metrics: bool = False


@dataclass(frozen=True)
class HealthResponse:
    status: StructuredStatus
    serving: bool
    state: str
    active_sessions: int = 0


@dataclass(frozen=True)
class SessionRequest:
    schema_version: int
    request_id: str
    work_id: str
    route_session: str
    route_epoch: int
    artifact_fingerprint: str
    runtime_fingerprint: str
    shard_start: int
    shard_end: int
    effective_start: int
    cache_expectation: CacheExpectation = CacheExpectation.NONE
    deadline_seconds: float = 30.0

    def validate(self) -> None:
        if self.schema_version not in SUPPORTED_VERSIONS:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.UNSUPPORTED_VERSION, "unsupported session schema version",
                details={"requested": str(self.schema_version), "supported": "1"},
            ))
        if not self.request_id or not self.work_id or not self.route_session:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.INVALID_ARGUMENT, "request, work, and route-session IDs are required",
            ))
        if self.route_epoch < 0 or self.shard_start < 0 or self.shard_end <= self.shard_start:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.INVALID_ARGUMENT, "invalid route epoch or shard range",
            ))
        if not self.shard_start <= self.effective_start <= self.shard_end:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.INVALID_ARGUMENT, "effective start must be inside the shard range",
            ))
        if self.deadline_seconds <= 0:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.INVALID_ARGUMENT, "session deadline must be positive",
            ))


@dataclass(frozen=True)
class SessionFrame:
    phase: SessionPhase
    position: int
    idempotency_step: int
    payload: Any
    cache_expectation: CacheExpectation = CacheExpectation.NONE
    size_bytes: int = 0

    def validate(self, limits: FlowControlLimits) -> None:
        if self.position < 0 or self.idempotency_step < 0:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.INVALID_ARGUMENT, "position and idempotency step must be non-negative",
            ))
        if self.size_bytes < 0 or self.size_bytes > limits.max_frame_bytes:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.RESOURCE_EXHAUSTED, "session frame exceeds max_frame_bytes",
            ))


@dataclass(frozen=True)
class SessionResult:
    status: StructuredStatus
    cache_result: CacheResult = CacheResult.NOT_REQUESTED
    position: int = 0
    idempotency_step: int = 0
    payload: Any = None


@dataclass(frozen=True)
class CancelRequest:
    schema_version: int
    request_id: str
    work_id: str
    route_session: str
    route_epoch: int
    reason: str = ""


@dataclass(frozen=True)
class ReleaseRequest:
    schema_version: int
    request_id: str
    work_id: str
    route_session: str
    route_epoch: int


@dataclass(frozen=True)
class MetricsRequest:
    schema_version: int = SCHEMA_VERSION


@dataclass(frozen=True)
class MetricsResponse:
    status: StructuredStatus
    active_sessions: int
    queued_frames: int
    inflight_bytes: int
    kv_entries: int
    generated_tokens: int
    cancelled_sessions: int


@dataclass
class SessionLifecycle:
    """Fail-closed state machine for one Route Session Activation Seam."""

    request: SessionRequest
    state: LifecycleState = LifecycleState.OPEN
    cancellation: CancellationToken = field(default_factory=CancellationToken)
    last_idempotency_step: int = -1
    _seen_steps: set[int] = field(default_factory=set, init=False, repr=False)

    _ALLOWED: ClassVar[dict[LifecycleState, frozenset[LifecycleState]]] = {
        LifecycleState.OPEN: frozenset({LifecycleState.PREFILLING, LifecycleState.CANCELLING, LifecycleState.RELEASING, LifecycleState.FAILED}),
        LifecycleState.PREFILLING: frozenset({LifecycleState.PREFILLING, LifecycleState.DECODING, LifecycleState.CANCELLING, LifecycleState.RELEASING, LifecycleState.FAILED}),
        LifecycleState.DECODING: frozenset({LifecycleState.DECODING, LifecycleState.CANCELLING, LifecycleState.RELEASING, LifecycleState.FAILED}),
        LifecycleState.CANCELLING: frozenset({LifecycleState.CANCELLED, LifecycleState.RELEASING, LifecycleState.FAILED}),
        LifecycleState.CANCELLED: frozenset({LifecycleState.RELEASING, LifecycleState.RELEASED}),
        LifecycleState.RELEASING: frozenset({LifecycleState.RELEASED, LifecycleState.FAILED}),
        LifecycleState.RELEASED: frozenset(),
        LifecycleState.FAILED: frozenset({LifecycleState.RELEASING, LifecycleState.RELEASED}),
    }

    def _transition(self, target: LifecycleState) -> None:
        if target not in self._ALLOWED[self.state]:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.MALFORMED_LIFECYCLE,
                f"cannot transition from {self.state.value} to {target.value}",
                details={"state": self.state.value, "target": target.value},
            ))
        self.state = target

    def apply(self, frame: SessionFrame) -> None:
        frame.validate(FlowControlLimits())
        if self.cancellation.cancelled and frame.phase is not SessionPhase.PREFILL:
            raise LifecycleContractError(StructuredStatus(StatusCode.CANCELLED, "session cancellation propagated"))
        if frame.idempotency_step in self._seen_steps:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.ALREADY_EXISTS, "duplicate idempotency step", details={"step": str(frame.idempotency_step)},
            ))
        if frame.idempotency_step <= self.last_idempotency_step:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.MALFORMED_LIFECYCLE, "idempotency steps must increase monotonically",
            ))
        target = LifecycleState.PREFILLING if frame.phase is SessionPhase.PREFILL else LifecycleState.DECODING
        self._transition(target)
        self._seen_steps.add(frame.idempotency_step)
        self.last_idempotency_step = frame.idempotency_step

    def cancel(self, reason: str = "") -> StructuredStatus:
        if self.state in {LifecycleState.RELEASED, LifecycleState.RELEASING}:
            return StructuredStatus(StatusCode.FAILED_PRECONDITION, "session is already releasing or released")
        if self.state is LifecycleState.CANCELLED:
            return StructuredStatus(StatusCode.OK, "session already cancelled")
        self._transition(LifecycleState.CANCELLING)
        self.cancellation.cancel()
        self._transition(LifecycleState.CANCELLED)
        return StructuredStatus(StatusCode.CANCELLED, reason or "session cancelled")

    def release(self) -> StructuredStatus:
        if self.state is LifecycleState.RELEASED:
            return StructuredStatus(StatusCode.OK, "session already released")
        self._transition(LifecycleState.RELEASING)
        self.cancellation.cancel()
        self._transition(LifecycleState.RELEASED)
        return StructuredStatus(StatusCode.OK, "session released")


@dataclass(frozen=True)
class ShardRpcContract:
    """Service/method and operational rules for generated gRPC bindings."""

    schema: str = SCHEMA_NAME
    version: int = SCHEMA_VERSION
    methods: tuple[RpcName, ...] = tuple(RpcName)
    deadlines: DeadlinePolicy = DeadlinePolicy()
    flow_control: FlowControlLimits = FlowControlLimits()
    tls_auth: TlsAuthHooks = TlsAuthHooks()

    def validate(self) -> None:
        if self.schema != SCHEMA_NAME or self.version not in SUPPORTED_VERSIONS:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.UNSUPPORTED_VERSION, "unsupported Shard RPC contract version",
            ))
        if self.methods != tuple(RpcName):
            raise LifecycleContractError(StructuredStatus(
                StatusCode.FAILED_PRECONDITION, "contract must expose the complete lifecycle RPC set",
            ))
        self.tls_auth.validate()

    def to_bytes(self) -> bytes:
        self.validate()
        data = {
            "schema": self.schema,
            "version": self.version,
            "methods": [method.value for method in self.methods],
            "deadlines": self.deadlines.__dict__,
            "flow_control": self.flow_control.__dict__,
            "tls_auth": self.tls_auth.__dict__,
        }
        return json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8")

    @classmethod
    def from_bytes(cls, raw: bytes) -> "ShardRpcContract":
        data = json.loads(raw)
        if int(data.get("version", -1)) not in SUPPORTED_VERSIONS:
            raise LifecycleContractError(StructuredStatus(
                StatusCode.UNSUPPORTED_VERSION, "unsupported Shard RPC contract version",
            ))
        methods = tuple(RpcName(item) for item in data.get("methods", ()))
        contract = cls(
            schema=str(data.get("schema", "")),
            version=int(data["version"]),
            methods=methods,
            deadlines=DeadlinePolicy(**data.get("deadlines", {})),
            flow_control=FlowControlLimits(**data.get("flow_control", {})),
            tls_auth=TlsAuthHooks(**data.get("tls_auth", {})),
        )
        contract.validate()
        return contract