from __future__ import annotations

import uuid
from datetime import datetime
from typing import Annotated, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator

from modelforge_api.domain.serving import (
    EmbeddingInvokeRequest,
    OCRInvokeRequest,
    RerankingInvokeRequest,
    SpeechTranscriptionInvokeRequest,
    VisionEmbeddingInvokeRequest,
)


class LabModel(BaseModel):
    model_config = ConfigDict(extra="forbid")


class LabTarget(LabModel):
    deployment_candidate_id: uuid.UUID
    artifact_set_id: uuid.UUID
    runtime_profile_id: uuid.UUID
    compute_node_id: uuid.UUID
    execution_approval_id: uuid.UUID


class LabEmbeddingInput(LabModel):
    capability: Literal["rag.embedding"]
    payload: EmbeddingInvokeRequest


class LabRerankingInput(LabModel):
    capability: Literal["rag.reranking"]
    payload: RerankingInvokeRequest


class LabOCRInput(LabModel):
    capability: Literal["document.ocr"]
    payload: OCRInvokeRequest


class LabVisionInput(LabModel):
    capability: Literal["vision.embedding"]
    payload: VisionEmbeddingInvokeRequest


class LabSpeechInput(LabModel):
    capability: Literal["speech.transcription"]
    payload: SpeechTranscriptionInvokeRequest


LabInput = Annotated[
    LabEmbeddingInput | LabRerankingInput | LabOCRInput | LabVisionInput | LabSpeechInput,
    Field(discriminator="capability"),
]


class LabLimits(LabModel):
    timeout_seconds: int = Field(default=60, ge=5, le=300, strict=True)
    queue_timeout_seconds: int = Field(default=10, ge=1, le=60, strict=True)
    retention_seconds: int = Field(default=300, ge=60, le=3600, strict=True)
    max_output_bytes: int = Field(default=65536, ge=1, le=1048576, strict=True)
    concurrency: Literal[1] = 1

    @model_validator(mode="after")
    def bounded_queue(self) -> LabLimits:
        if self.queue_timeout_seconds > self.timeout_seconds:
            raise ValueError("queue timeout cannot exceed total timeout")
        return self


class LabPreflightRequest(LabModel):
    target: LabTarget
    input: LabInput
    contract_version: Literal[1] = 1
    limits: LabLimits = Field(default_factory=LabLimits)


class LabExecutionStartRequest(LabPreflightRequest):
    """An explicitly selected embedding experiment; body input is never echoed in status."""

    deployment_id: uuid.UUID
    experiment_route_id: uuid.UUID


class LabInputSummary(LabModel):
    sha256: str
    item_count: int
    encoded_bytes: int


class LabBinding(LabTarget):
    compatibility_assessment_id: uuid.UUID
    runtime_probe_id: uuid.UUID
    capability: str
    contract_version: Literal[1] = 1
    runtime_profile_fingerprint: str
    assessment_fingerprint: str
    approval_evidence_fingerprint: str
    probe_environment_fingerprint: str
    accelerator_id: uuid.UUID | None = None


class LabPreflightResponse(LabModel):
    input_summary: LabInputSummary
    binding: LabBinding | None = None
    limits: LabLimits
    blockers: list[str]
    checks_passed: bool
    execution_started: Literal[False] = False
    permission_granted: Literal[False] = False


LabExecutionStatus = Literal[
    "accepted",
    "preparing",
    "queued",
    "running",
    "cancel_requested",
    "completed",
    "failed",
    "cancelled",
    "interrupted",
]


class LabExecutionRunResponse(LabModel):
    id: uuid.UUID
    capability: Literal["rag.embedding@1"] = "rag.embedding@1"
    status: LabExecutionStatus
    target_snapshot: dict[str, object]
    input_summary: LabInputSummary
    output_metadata: dict[str, object] | None = None
    worker_observations: dict[str, object] | None = None
    failure_code: str | None = None
    input_payload_expires_at: datetime
    retention_seconds: int = Field(ge=60, le=3600)
    created_at: datetime
    started_at: datetime | None = None
    cancellation_requested_at: datetime | None = None
    cancellation_acknowledged_at: datetime | None = None
    finished_at: datetime | None = None