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)
    policy_revision: int = Field(default=0, ge=0)
    planner_transport: Literal["direct", "cloud"] = "direct"


class HostGovernancePolicyModel(BaseModel):
    revision: int = Field(ge=1)
    self_submission_enabled: bool
    max_active_tasks: int | None = None
    daily_token_budget: int | None = None


class HeartbeatResponse(BaseModel):
    host_id: str
    accepted_devices: int
    received_at: datetime
    policy_revision: int = Field(default=0, ge=0)
    policy: HostGovernancePolicyModel | None = None


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 TaskProgressModel(BaseModel):
    step_index: int = Field(ge=0)
    step_status: Literal["running", "completed", "failed"]
    summary: str = Field(default="", max_length=2000)


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)
    progress: TaskProgressModel | None = None


class LeaseRenewalResponse(BaseModel):
    status: Literal["renewed"]
    lease_expires_at: datetime
    cancel_requested: bool = False


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", "cancelled"]
    failure_reason: str | None = None
    result: dict[str, Any] | None = None


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


class HostTaskSubmissionRequest(BaseModel):
    host_id: str = Field(min_length=1)
    goal: str = Field(min_length=1)
    device_id: str | None = Field(default=None, min_length=1)


class HostTaskSubmissionResponse(BaseModel):
    task_id: str


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)
    task_id: str | None = Field(default=None, min_length=1)
    attempt: int | None = Field(default=None, ge=1)
    lease_id: str | None = Field(default=None, min_length=1)


class PlannerDecisionResponse(BaseModel):
    tool_name: str
    arguments: dict[str, Any] = Field(default_factory=dict)
    input_tokens: int | None = Field(default=None, ge=0)
    output_tokens: int | None = Field(default=None, ge=0)
    total_tokens: int | None = Field(default=None, ge=0)


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