from __future__ import annotations

from datetime import datetime
from typing import Any, Literal

from pydantic import BaseModel, Field


class HostEnrollmentRequest(BaseModel):
    agent_instance_id: str = Field(min_length=1, max_length=256)
    host_token: str = Field(min_length=32, max_length=512)
    display_name: str | None = Field(default=None, max_length=256)


class HostEnrollmentResponse(BaseModel):
    host_id: str


class DeviceEnrollmentRequest(BaseModel):
    local_device_id: str = Field(min_length=1, max_length=256)
    driver_type: str = Field(min_length=1, max_length=128)
    name: str | None = Field(default=None, max_length=256)
    capability_tags: list[str] = Field(default_factory=list)


class DeviceEnrollmentResponse(BaseModel):
    device_id: str


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"]


class StaleLeaseConflict(BaseModel):
    code: Literal["stale_lease"] = "stale_lease"
    detail: str


class PlannerToolSpecModel(BaseModel):
    name: str = Field(min_length=1)
    description: str = ""
    parameters: dict[str, Any] = Field(default_factory=dict)


class PlannerDecisionRequest(BaseModel):
    host_id: str = Field(min_length=1)
    system_prompt: str
    user_prompt: str
    screenshot_base64: str | None = None
    tools: list[PlannerToolSpecModel] = Field(default_factory=list)
    timeout_seconds: float = Field(default=30.0, gt=0, le=120)


class PlannerDecisionResponse(BaseModel):
    tool_name: str
    arguments: dict[str, Any] = Field(default_factory=dict)


class PlannerDecisionError(BaseModel):
    code: Literal["planner_unavailable"] = "planner_unavailable"
    detail: str