from __future__ import annotations

import base64
import binascii
import math
import re
import uuid
from datetime import datetime
from typing import Any, Literal

from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, field_validator, model_validator

CAPABILITY_VERSION_KEY = re.compile(r"^[a-z][a-z0-9_]*(?:\.[a-z][a-z0-9_]*)+@[1-9][0-9]*$")
ServingOperation = Literal["load", "invoke", "health", "drain", "unload"]


class ServingModel(BaseModel):
    model_config = ConfigDict(extra="forbid", from_attributes=True)


class SupplyChainReview(ServingModel):
    exact_revision_reviewed: Literal[True]
    artifact_hashes_reviewed: Literal[True]
    safetensors_only: Literal[True]
    remote_code_required: Literal[False]
    pickle_present: Literal[False]
    scanner_evidence_reviewed: Literal[True]
    provenance_complete: Literal[True]
    runtime_offline_reviewed: Literal[True]
    dependency_provenance_reviewed: Literal[True]
    license_reviewed: Literal[True]
    license_identifier: str = Field(min_length=1, max_length=255)


class ProductionApprovalCreate(ServingModel):
    capability: Literal["rag.embedding"] = "rag.embedding"
    contract_version: Literal[1] = 1
    approved_by: str = Field(min_length=1, max_length=255)
    reason: str = Field(min_length=16, max_length=4000)
    supply_chain_review: SupplyChainReview
    deployment_config: dict[str, Any] = Field(default_factory=dict)


class ProductionApprovalResponse(ServingModel):
    id: uuid.UUID
    deployment_candidate_id: uuid.UUID
    capability_contract_id: uuid.UUID
    artifact_set_id: uuid.UUID
    runtime_profile_id: uuid.UUID
    compute_node_id: uuid.UUID
    deployment_config: dict[str, Any]
    supply_chain_evidence: dict[str, Any]
    evidence_fingerprint: str
    status: str
    approved_by: str
    reason: str
    approved_at: datetime
    revoked_at: datetime | None
    stale: bool = False


class CapabilityPromotionCreate(ServingModel):
    production_approval_id: uuid.UUID
    residency_policy: Literal["always_warm", "keep_warm", "load_on_demand"] = "keep_warm"
    keep_warm_seconds: int = Field(default=900, ge=5, le=86400)
    max_concurrency: Literal[1] = 1
    max_queue_depth: int = Field(default=16, ge=1, le=1024)
    routing_weight: int = Field(default=100, ge=0, le=100)
    rollback_policy: dict[str, Any] = Field(
        default_factory=lambda: {"mode": "drain_to_unavailable", "previous_deployment_id": None}
    )


class EmbeddingSpaceResponse(ServingModel):
    id: uuid.UUID
    identity_digest: str
    dimension: int
    normalized: bool
    migration_class: Literal["requires_reindex"]
    identity_facts: dict[str, Any]
    created_at: datetime


class ResourceEnvelopeResponse(ServingModel):
    id: uuid.UUID
    runtime_probe_id: uuid.UUID
    accelerator_kind: str
    accelerator_uuid: str
    environment_fingerprint: str
    concurrency: int
    batch_size: int
    max_sequence_length: int
    baseline_vram_bytes: int
    resident_vram_bytes: int
    peak_vram_bytes: int
    required_vram_bytes: int
    cold_load_time_ms: float
    inference_latency_ms: float
    stale: bool
    stale_reason: str | None


class ResidencyResponse(ServingModel):
    id: uuid.UUID
    state: str
    worker_instance_id: str | None
    load_count: int
    active_requests: int
    measured_resident_vram_bytes: int
    external_baseline_vram_bytes: int
    health: dict[str, Any]
    resident_since: datetime | None
    last_used_at: datetime | None
    failure_code: str | None
    failure_message: str | None
    generation: int = 1
    transition_reason: str | None = None


class CapabilityDeploymentResponse(ServingModel):
    id: uuid.UUID
    route_generation: int = Field(ge=0)
    capability: str
    contract_version: int
    deployment_candidate_id: uuid.UUID
    production_approval_id: uuid.UUID | None
    execution_approval_id: uuid.UUID | None
    artifact_set_id: uuid.UUID
    runtime_profile_id: uuid.UUID
    compute_node_id: uuid.UUID
    accelerator_id: uuid.UUID
    embedding_space: EmbeddingSpaceResponse | None
    resource_envelope: ResourceEnvelopeResponse
    residency: ResidencyResponse
    channel: str
    status: str
    production: bool
    health_status: str
    routing_weight: int
    fallback_policy: dict[str, Any]
    residency_policy: str
    keep_warm_seconds: int
    max_concurrency: int
    max_queue_depth: int
    config_fingerprint: str
    provenance: dict[str, Any]
    rollback_policy: dict[str, Any]
    promoted_at: datetime | None
    created_at: datetime


class ServiceClientCreate(ServingModel):
    name: str = Field(min_length=1, max_length=255, pattern=r"^[a-z][a-z0-9-]*$")
    allowed_capabilities: list[str] = Field(min_length=1, max_length=32)
    requests_per_minute: int = Field(default=60, ge=1, le=10000)
    max_concurrent_requests: int = Field(default=1, ge=1, le=64)
    workload_priority: Literal["production", "interactive", "background", "benchmark"] = (
        "production"
    )
    project_key: str | None = Field(default=None, pattern=r"^[a-z][a-z0-9-]*$")
    integration_environment: Literal["production", "shadow", "evaluation", "lab"] | None = None
    purpose: str | None = Field(default=None, min_length=3, max_length=500)
    credential_expires_at: AwareDatetime | None = None
    project_binding_id: uuid.UUID | None = None

    @field_validator("allowed_capabilities")
    @classmethod
    def validate_capability_scopes(cls, value: list[str]) -> list[str]:
        if len(set(value)) != len(value) or any(
            not CAPABILITY_VERSION_KEY.fullmatch(v) for v in value
        ):
            raise ValueError("capability scopes must be unique versioned capability keys")
        return value

    @model_validator(mode="after")
    def validate_project_scope(self) -> ServiceClientCreate:
        project_fields = (self.project_key, self.integration_environment, self.purpose)
        if any(value is not None for value in project_fields) and not all(
            value is not None for value in project_fields
        ):
            raise ValueError(
                "project_key, integration_environment and purpose are required together"
            )
        if self.project_key is not None and len(self.allowed_capabilities) != 1:
            raise ValueError("a project service identity must bind exactly one capability")
        if self.project_binding_id is not None and self.project_key is None:
            raise ValueError("an explicit binding requires a project identity")
        return self


class ServiceClientResponse(ServingModel):
    id: uuid.UUID
    name: str
    status: str
    allowed_capabilities: list[str]
    requests_per_minute: int
    max_concurrent_requests: int
    workload_priority: str
    project_key: str | None
    project_binding_id: uuid.UUID | None
    integration_environment: str | None
    purpose: str | None
    credential_prefix: str | None
    credential_id: uuid.UUID | None = None
    credential_created_at: datetime | None
    credential_expires_at: datetime | None
    credential_revoked_at: datetime | None
    last_used_at: datetime | None
    created_at: datetime


class ServiceClientCreated(ServiceClientResponse):
    credential: str


class ServiceCredentialRotate(ServingModel):
    expected_credential_id: uuid.UUID
    overlap_seconds: int = Field(default=0, strict=True, ge=0, le=86400)
    expires_at: AwareDatetime


class ServiceCredentialRevoke(ServingModel):
    expected_credential_id: uuid.UUID
    confirmation: str = Field(min_length=1, max_length=255)


class ServiceCredentialResponse(ServingModel):
    id: uuid.UUID
    secret_prefix: str
    created_at: datetime
    expires_at: datetime | None
    superseded_at: datetime | None
    revoked_at: datetime | None
    last_used_at: datetime | None


class ProjectFitEvidenceCreate(ServingModel):
    project_key: str = Field(pattern=r"^[a-z][a-z0-9-]*$")
    capability: str = Field(pattern=r"^[a-z][a-z0-9.]*@[1-9][0-9]*$")
    environment: Literal["production", "shadow", "evaluation", "lab"]
    recommendation: Literal["PROMOTION_ELIGIBLE", "KEEP_LAB", "BLOCKED", "REQUIRES_MORE_EVIDENCE"]
    evidence_class: Literal["OWNER_PHOTO", "PUBLIC_PHYSICAL_CAPTURE", "CATALOG_REFERENCE"]
    engineering_integration: Literal["PASS", "INCOMPLETE", "BLOCKED"]
    production_validation: Literal["DEFERRED_EXTERNAL_VALIDATION", "REQUIRED", "SATISFIED"]
    production_action: Literal["NONE"] = "NONE"
    case_count: int = Field(ge=0, le=1_000_000)
    metric_values: dict[str, float | int | str | bool | None]
    critical_errors: int = Field(ge=0)
    blockers: list[str] = Field(default_factory=list, max_length=100)
    evidence: dict[str, Any] = Field(default_factory=dict)
    capability_deployment_id: uuid.UUID | None = None


class ProjectFitEvidenceResponse(ServingModel):
    id: uuid.UUID
    project_key: str
    capability: str
    environment: str
    recommendation: str
    evidence_class: str
    engineering_integration: str
    production_validation: str
    production_action: str
    case_count: int
    metric_values: dict[str, Any]
    critical_errors: int
    blockers: list[str]
    evidence_digest: str
    evidence: dict[str, Any]
    capability_deployment_id: uuid.UUID | None
    created_at: datetime


class ProjectIntegrationUsage(ServingModel):
    request_volume: int
    successful_requests: int
    error_count: int
    last_used_at: datetime | None
    latency_p50_ms: float | None
    latency_p95_ms: float | None


class ProjectIntegrationResponse(ServingModel):
    project_key: str
    project_name: str
    capability: str
    environment: str
    state: str
    purpose: str
    client_id: uuid.UUID
    client_name: str
    deployment_id: uuid.UUID | None
    resource_impact_bytes: int | None
    usage: ProjectIntegrationUsage
    project_fit: ProjectFitEvidenceResponse | None


class EmbeddingInvokeRequest(ServingModel):
    input: str | list[str]
    input_type: Literal["raw", "query", "document"] = "raw"

    @model_validator(mode="after")
    def validate_input(self) -> EmbeddingInvokeRequest:
        values = [self.input] if isinstance(self.input, str) else self.input
        if not values or any(not item.strip() for item in values):
            raise ValueError("input must contain at least one non-empty string")
        return self

    def inputs(self) -> list[str]:
        return [self.input] if isinstance(self.input, str) else self.input


class CapabilityExperimentCreate(ServingModel):
    route_key: str = Field(min_length=3, max_length=128, pattern=r"^[a-z][a-z0-9-]+$")
    execution_approval_id: uuid.UUID
    capability: Literal[
        "rag.embedding",
        "rag.reranking",
        "document.ocr",
        "vision.embedding",
        "speech.transcription",
    ] = "rag.embedding"
    purpose: str = Field(min_length=16, max_length=2000)
    expected_dimension: int = Field(default=1024, ge=32, le=8192)
    residency_policy: Literal["keep_warm", "load_on_demand"] = "keep_warm"
    keep_warm_seconds: int = Field(default=300, ge=5, le=86400)
    max_queue_depth: int = Field(default=16, ge=1, le=1024)


class CapabilityExperimentResponse(ServingModel):
    id: uuid.UUID
    route_key: str
    capability: str
    contract_version: int
    capability_deployment_id: uuid.UUID
    status: str
    purpose: str
    evidence: dict[str, Any]
    deployment: CapabilityDeploymentResponse
    created_at: datetime


class GatewayTiming(ServingModel):
    validation_ms: float = 0.0
    resolution_ms: float = 0.0
    scheduling_ms: float = 0.0
    payload_ms: float = 0.0
    dispatch_ms: float = 0.0
    worker_preprocess_ms: float = 0.0
    worker_serialize_ms: float = 0.0
    worker_total_ms: float = 0.0
    completion_transport_ms: float = 0.0
    result_ms: float = 0.0
    gateway_ms: float
    # Backwards-compatible alias for the scheduler/lease wait used by M5/M6.
    queue_ms: float
    load_ms: float
    inference_ms: float
    total_ms: float


class GatewayExecution(ServingModel):
    cold: bool
    node: str
    residency: str
    load_count: int
    timings: GatewayTiming


class GatewayUsage(ServingModel):
    input_count: int
    input_tokens: int


class EmbeddingInvokeResponse(ServingModel):
    capability: Literal["rag.embedding@1"] = "rag.embedding@1"
    dimension: int = Field(ge=1, le=65_536)
    normalized: Literal[True] = True
    embedding_space_id: uuid.UUID
    data: list[list[float]]
    request_id: uuid.UUID
    execution: GatewayExecution
    usage: GatewayUsage


class StableEmbeddingInvokeResponse(EmbeddingInvokeResponse):
    """Public stable contract; lab experiment routes may use another dimension."""

    dimension: Literal[1024] = 1024


class RerankDocument(ServingModel):
    id: str = Field(min_length=1, max_length=255)
    text: str = Field(min_length=1, max_length=16_384)


class RerankingInvokeRequest(ServingModel):
    query: str = Field(min_length=1, max_length=4_000)
    documents: list[RerankDocument] = Field(min_length=1, max_length=40)
    top_n: int = Field(default=10, ge=1, le=10)

    @model_validator(mode="after")
    def validate_reranking_input(self) -> RerankingInvokeRequest:
        if not self.query.strip() or any(not document.text.strip() for document in self.documents):
            raise ValueError("query and document text must be non-empty")
        if len({document.id for document in self.documents}) != len(self.documents):
            raise ValueError("document ids must be unique")
        if self.top_n > len(self.documents):
            raise ValueError("top_n cannot exceed the document count")
        if len(self.query) + sum(len(document.text) for document in self.documents) > 262_144:
            raise ValueError("rerank request is too large")
        return self


class RerankingResult(ServingModel):
    id: str
    score: float
    rank: int = Field(ge=1, le=10)

    @field_validator("score")
    @classmethod
    def finite_score(cls, value: float) -> float:
        if not math.isfinite(value):
            raise ValueError("reranking score must be finite")
        return value


class RerankingInvokeResponse(ServingModel):
    capability: Literal["rag.reranking@1"] = "rag.reranking@1"
    results: list[RerankingResult]
    request_id: uuid.UUID
    execution: GatewayExecution
    usage: GatewayUsage


def _bounded_base64(value: str, maximum_bytes: int) -> str:
    try:
        decoded = base64.b64decode(value, validate=True)
    except (binascii.Error, ValueError) as exc:
        raise ValueError("content must be canonical base64") from exc
    if not decoded or len(decoded) > maximum_bytes:
        raise ValueError(f"decoded content must be between 1 and {maximum_bytes} bytes")
    return value


class OCRInvokeRequest(ServingModel):
    content_base64: str
    media_type: Literal["image/png", "image/jpeg"]
    language_hint: Literal["nl", "en", "auto"] = "auto"

    @field_validator("content_base64")
    @classmethod
    def bounded_content(cls, value: str) -> str:
        return _bounded_base64(value, 8 * 1024 * 1024)


class OCRBlock(ServingModel):
    text: str
    order: int = Field(ge=0)
    bbox: list[int] | None = Field(default=None, min_length=4, max_length=4)
    confidence: float | None = Field(default=None, ge=0, le=1)


class OCRPage(ServingModel):
    page: Literal[1] = 1
    width: int = Field(ge=1, le=4096)
    height: int = Field(ge=1, le=4096)
    blocks: list[OCRBlock]


class OCRInvokeResponse(ServingModel):
    capability: Literal["document.ocr@1"] = "document.ocr@1"
    text: str
    pages: list[OCRPage] = Field(min_length=1, max_length=1)
    confidence: float | None = Field(default=None, ge=0, le=1)
    request_id: uuid.UUID
    execution: GatewayExecution


class VisionEmbeddingItem(ServingModel):
    image_base64: str | None = None
    media_type: Literal["image/png", "image/jpeg"] | None = None
    text: str | None = Field(default=None, min_length=1, max_length=4096)

    @model_validator(mode="after")
    def exactly_one_modality(self) -> VisionEmbeddingItem:
        image = self.image_base64 is not None
        text = self.text is not None
        if image == text or image != (self.media_type is not None):
            raise ValueError("each item must contain exactly one typed image or text")
        if self.image_base64 is not None:
            _bounded_base64(self.image_base64, 8 * 1024 * 1024)
        return self


class VisionEmbeddingInvokeRequest(ServingModel):
    items: list[VisionEmbeddingItem] = Field(min_length=1, max_length=4)

    @model_validator(mode="after")
    def bounded_total(self) -> VisionEmbeddingInvokeRequest:
        encoded = sum(len(item.image_base64 or "") for item in self.items)
        if encoded > 11_184_812:
            raise ValueError("combined image payload exceeds 8 MiB")
        return self


class VisionEmbeddingInvokeResponse(ServingModel):
    capability: Literal["vision.embedding@1"] = "vision.embedding@1"
    dimension: int = Field(ge=32, le=8192)
    normalized: Literal[True] = True
    embedding_space_id: uuid.UUID
    data: list[list[float]] = Field(min_length=1, max_length=4)
    request_id: uuid.UUID
    execution: GatewayExecution


class SpeechTranscriptionInvokeRequest(ServingModel):
    audio_base64: str
    media_type: Literal["audio/wav"] = "audio/wav"
    language: Literal["nl", "en", "fr", "auto"] = "auto"

    @field_validator("audio_base64")
    @classmethod
    def bounded_audio(cls, value: str) -> str:
        return _bounded_base64(value, 16 * 1024 * 1024)


class TranscriptionSegment(ServingModel):
    start_seconds: float = Field(ge=0)
    end_seconds: float = Field(ge=0)
    text: str


class SpeechTranscriptionInvokeResponse(ServingModel):
    capability: Literal["speech.transcription@1"] = "speech.transcription@1"
    text: str
    language: str
    duration_seconds: float = Field(gt=0, le=120)
    segments: list[TranscriptionSegment] = Field(default_factory=list)
    request_id: uuid.UUID
    execution: GatewayExecution


class OpenAIEmbeddingRequest(ServingModel):
    model: Literal["rag.embedding", "rag.embedding@1"]
    input: str | list[str]


class OpenAIEmbeddingItem(ServingModel):
    object: Literal["embedding"] = "embedding"
    index: int
    embedding: list[float]


class OpenAIUsage(ServingModel):
    prompt_tokens: int
    total_tokens: int


class OpenAIEmbeddingResponse(ServingModel):
    object: Literal["list"] = "list"
    data: list[OpenAIEmbeddingItem]
    model: Literal["rag.embedding"] = "rag.embedding"
    usage: OpenAIUsage


class SchedulerBudgetResponse(ServingModel):
    compute_node_id: uuid.UUID
    node_name: str
    accelerator_id: uuid.UUID
    accelerator_uuid: str
    accelerator_name: str
    total_vram_bytes: int
    observed_used_vram_bytes: int
    external_vram_bytes: int
    resident_vram_bytes: int
    leased_vram_bytes: int
    safety_reserve_bytes: int
    schedulable_free_vram_bytes: int
    pressure: bool
    pressure_state: Literal["NORMAL", "ELEVATED", "HIGH", "CRITICAL"] = "NORMAL"
    attribution_confidence: Literal["KNOWN", "ESTIMATED", "UNKNOWN"] = "UNKNOWN"
    invariant_delta_bytes: int = 0
    policy_revision: str = "m10-v1"
    observed_at: datetime | None


class PlacementPlanRequest(ServingModel):
    priority: Literal["production", "interactive", "background", "lab"] = "interactive"
    deadline_remaining_ms: float | None = Field(default=None, gt=0, le=900_000)


class PlacementEvictionResponse(ServingModel):
    deployment_id: uuid.UUID
    capability: str
    expected_reclaimed_bytes: int
    reason: str


class PlacementPlanResponse(ServingModel):
    id: uuid.UUID | None = None
    deployment_id: uuid.UUID
    capability: str
    node_id: uuid.UUID
    accelerator_id: uuid.UUID
    policy_revision: str
    verdict: Literal[
        "ADMIT",
        "ADMIT_AFTER_EVICTION",
        "QUEUE",
        "REJECT_CAPACITY",
        "REJECT_HEALTH",
        "REJECT_POLICY",
    ]
    reason_codes: list[str]
    required_vram_bytes: int
    headroom_before_bytes: int
    headroom_after_bytes: int
    expected_cold_load_ms: float
    evictions: list[PlacementEvictionResponse] = Field(default_factory=list)
    decision_fingerprint: str
    dry_run: bool
    created_at: datetime


class SchedulerPolicyResponse(ServingModel):
    revision: str
    active: bool
    configuration: dict[str, Any]
    created_at: datetime


class SchedulerMetricsResponse(ServingModel):
    counters: dict[str, int]
    gauges: dict[str, int | str]
    queue_seconds: dict[str, float]


class SchedulerPolicyUpdate(ServingModel):
    lab_paused: bool


class CoResidencyEvidenceResponse(ServingModel):
    left_deployment_id: uuid.UUID
    left_capability: str
    right_deployment_id: uuid.UUID
    right_capability: str
    left_alone_bytes: int
    right_alone_bytes: int
    expected_combined_bytes: int
    measured_combined_bytes: int | None
    status: Literal["PROVEN_SAFE", "EXPECTED_SAFE", "NOT_SAFE", "UNKNOWN"]
    evidence: dict[str, Any] = Field(default_factory=dict)
    measured_at: datetime | None


class CoResidencyEvidenceCreate(ServingModel):
    left_deployment_id: uuid.UUID
    right_deployment_id: uuid.UUID

    @model_validator(mode="after")
    def distinct_deployments(self) -> CoResidencyEvidenceCreate:
        if self.left_deployment_id == self.right_deployment_id:
            raise ValueError("co-residency evidence requires two distinct deployments")
        return self


class ResidencyPolicyUpdate(ServingModel):
    residency_policy: Literal["always_warm", "keep_warm", "load_on_demand", "lab_only"]
    keep_warm_seconds: int = Field(default=900, ge=5, le=86400)


class GatewayRequestResponse(ServingModel):
    request_id: uuid.UUID
    capability: str
    client: str | None
    status: str
    cold: bool | None
    queue_time_ms: float | None
    load_time_ms: float | None
    inference_time_ms: float | None
    total_latency_ms: float | None
    failure_code: str | None
    latency_breakdown: dict[str, float] = Field(default_factory=dict)
    created_at: datetime


class AgentServingJobLease(ServingModel):
    job_id: uuid.UUID
    lease_token: str = Field(min_length=32, max_length=512)
    lease_expires_at: datetime
    operation: ServingOperation
    deployment_id: uuid.UUID
    artifact_set_id: uuid.UUID
    revision_sha: str = Field(pattern=r"^[a-f0-9]{40,64}$")
    artifact_root: str
    artifact_relative_path: str
    expected_manifest: dict[str, Any]
    runtime_profile: dict[str, Any]
    runtime_environment: dict[str, Any]
    capability: Literal[
        "rag.embedding",
        "rag.reranking",
        "document.ocr",
        "vision.embedding",
        "speech.transcription",
    ] = "rag.embedding"
    embedding_space_id: uuid.UUID | None = None
    expected_dimension: int | None = Field(default=None, ge=32, le=8192)
    normalize: Literal[True] | None = True
    input: list[str] | None = None
    input_type: Literal["raw", "query", "document"] = "raw"
    rerank_query: str | None = None
    rerank_documents: list[RerankDocument] | None = None
    top_n: int | None = Field(default=None, ge=1, le=10)
    modality_payload: dict[str, Any] | None = None
    request_id: uuid.UUID | None = None
    lab_run_id: uuid.UUID | None = None

    @model_validator(mode="after")
    def validate_capability_payload(self) -> AgentServingJobLease:
        if self.operation != "invoke":
            return self
        if self.capability == "rag.embedding":
            if self.embedding_space_id is None or self.expected_dimension is None:
                raise ValueError("embedding invoke requires embedding identity")
            if self.input is None or self.rerank_query is not None:
                raise ValueError("embedding invoke payload is invalid")
        elif self.capability == "rag.reranking" and (
            self.rerank_query is None
            or self.rerank_documents is None
            or self.top_n is None
            or self.input is not None
        ):
            raise ValueError("reranking invoke payload is invalid")
        elif self.capability not in {"rag.embedding", "rag.reranking"} and (
            self.modality_payload is None or self.input is not None or self.rerank_query is not None
        ):
            raise ValueError("modality invoke requires a typed modality payload")
        return self


class AgentServingJobComplete(ServingModel):
    lease_token: str = Field(min_length=32, max_length=512)
    worker_instance_id: str = Field(min_length=1, max_length=255)
    state: str
    result: dict[str, Any] = Field(default_factory=dict)
    metrics: dict[str, Any] = Field(default_factory=dict)
    health: dict[str, Any] = Field(default_factory=dict)
    runtime_facts: dict[str, Any] = Field(default_factory=dict)


class AgentServingJobFailure(ServingModel):
    lease_token: str = Field(min_length=32, max_length=512)
    worker_instance_id: str = Field(min_length=1, max_length=255)
    failure_code: str = Field(min_length=1, max_length=64)
    failure_message: str = Field(min_length=1, max_length=2000)
    state: str = "failed"
    details: dict[str, Any] = Field(default_factory=dict)


class AgentServingJobCancellationCheck(ServingModel):
    lease_token: str = Field(min_length=32, max_length=512)


class AgentServingJobCancellationStatus(ServingModel):
    cancel_requested: bool


class AgentServingJobCancellationAck(ServingModel):
    lease_token: str = Field(min_length=32, max_length=512)
    worker_instance_id: str = Field(min_length=1, max_length=255)
    state: Literal["cold", "warm", "busy", "loading", "draining", "unloading", "failed"]


class AgentResidentState(ServingModel):
    deployment_id: uuid.UUID
    artifact_set_id: uuid.UUID
    runtime_profile_id: uuid.UUID
    runtime_environment_fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$")
    revision_sha: str = Field(pattern=r"^[a-f0-9]{40,64}$")
    state: Literal["loading", "warm", "busy", "draining", "unloading", "failed"]
    load_count: int = Field(ge=1)
    resident_vram_bytes: int = Field(ge=0)
    health: dict[str, Any] = Field(default_factory=dict)


class AgentServingStateReport(ServingModel):
    worker_instance_id: str = Field(min_length=1, max_length=255)
    deployment_id: uuid.UUID | None
    state: Literal["cold", "loading", "warm", "busy", "draining", "unloading", "failed"]
    load_count: int = Field(ge=0)
    health: dict[str, Any] = Field(default_factory=dict)
    observed_at: datetime
    generation: int = Field(default=1, ge=1)
    residencies: list[AgentResidentState] = Field(default_factory=list, max_length=32)


class AgentServingStateAck(ServingModel):
    accepted: bool = True
    desired_operation: Literal["none", "unload"] = "none"