from __future__ import annotations

from datetime import datetime
from typing import Any, Literal

from pydantic import BaseModel, Field


class DeviceSnapshotModel(BaseModel):
    device_id: str = Field(min_length=1)
    driver_type: str = Field(min_length=1)
    status: Literal["idle", "busy", "offline", "error"]
    capability_tags: list[str] = Field(default_factory=list)


class HeartbeatRequest(BaseModel):
    host_id: str = Field(min_length=1)
    address: str | None = None
    devices: list[DeviceSnapshotModel] = Field(default_factory=list)


class HeartbeatResponse(BaseModel):
    host_id: str
    accepted_devices: int
    received_at: datetime


class ClaimRequest(BaseModel):
    host_id: str = Field(min_length=1)
    timeout_seconds: float = Field(default=20.0, ge=0, le=60)


class AssignmentModel(BaseModel):
    task_id: str
    attempt: int = Field(ge=1)
    lease_id: str
    lease_expires_at: datetime
    host_id: str
    device_id: str
    goal: str | None = None
    workflow_definition_id: str | None = None


class ClaimResponse(BaseModel):
    assignment: AssignmentModel | None = None
    timed_out: bool = False


class LeaseRenewalRequest(BaseModel):
    host_id: str = Field(min_length=1)
    task_id: str = Field(min_length=1)
    attempt: int = Field(ge=1)
    lease_id: str = Field(min_length=1)


class LeaseRenewalResponse(BaseModel):
    status: Literal["renewed"]
    lease_expires_at: datetime


class TerminalResultRequest(BaseModel):
    host_id: str = Field(min_length=1)
    task_id: str = Field(min_length=1)
    attempt: int = Field(ge=1)
    lease_id: str = Field(min_length=1)
    status: Literal["done", "failed"]
    failure_reason: str | None = None
    result: dict[str, Any] | None = None


class TerminalResultResponse(BaseModel):
    status: Literal["recorded", "already_recorded"]