"""Shared models for HiveCompute coordinator and workers."""

from __future__ import annotations

from datetime import UTC, datetime
from enum import StrEnum

from pydantic import BaseModel, Field

# --- Enums ---


class JobStatus(StrEnum):
    PENDING = "pending"
    CLAIMED = "claimed"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"


class WorkerStatus(StrEnum):
    IDLE = "idle"
    BUSY = "busy"
    PAUSED = "paused"
    OFFLINE = "offline"


class JobQuality(StrEnum):
    FAST = "fast"
    BALANCED = "balanced"
    BEST = "best"


# --- Internal models ---


class ComputeJob(BaseModel):
    """A compute job as seen by the worker daemon."""

    id: str
    job_type: str = "transcribe"
    status: JobStatus = JobStatus.PENDING

    input_url: str | None = None
    input_data: dict = Field(default_factory=dict)

    timeout_seconds: int = 3600
    max_attempts: int = 3
    attempts: int = 0

    # Shard info (populated for sharded jobs)
    shard_id: str | None = None
    start_seconds: float | None = None
    end_seconds: float | None = None

    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))


class WorkerInfo(BaseModel):
    """Worker hardware and availability info."""

    worker_id: str
    status: WorkerStatus = WorkerStatus.OFFLINE
    available_handlers: list[str] = Field(default_factory=list)
    available_models: list[str] = Field(default_factory=list)
    model_capabilities: list[dict] = Field(default_factory=list)

    chip: str | None = None
    memory_gb: int | None = None
    gpu_cores: int | None = None

    last_heartbeat: datetime | None = None
    jobs_completed: int = 0


# --- API request/response models ---


class JobSubmitRequest(BaseModel):
    audio_url: str
    language: str = "en"
    quality: JobQuality = JobQuality.BALANCED
    model: str | None = None
    priority: int = Field(default=3, ge=1, le=5)
    callback_url: str | None = None
    audio_duration: float | None = None  # Seconds — if provided, enables sharding for long audio


class JobSubmitResponse(BaseModel):
    job_id: str
    status: str = "pending"
    created_at: str
    resolved_model: str


class JobListItemResponse(BaseModel):
    job_id: str
    job_type: str
    status: str
    priority: int
    claimed_by: str | None = None
    created_at: str
    completed_at: str | None = None
    attempts: int = 0


class JobListResponse(BaseModel):
    jobs: list[JobListItemResponse]
    total: int
    limit: int
    offset: int


class JobStatusResponse(BaseModel):
    job_id: str
    job_type: str
    status: str
    priority: int
    claimed_by: str | None = None
    created_at: str
    claimed_at: str | None = None
    completed_at: str | None = None
    attempts: int = 0
    max_attempts: int = 3


class JobResultResponse(BaseModel):
    job_id: str
    status: str
    result: dict | None = None
    completed_at: str | None = None
    message: str | None = None


class WorkerRegisterRequest(BaseModel):
    worker_id: str = Field(min_length=1, max_length=128, pattern=r"^[a-zA-Z0-9._-]+$")
    chip: str | None = None
    memory_gb: int | None = None
    gpu_cores: int | None = None
    available_handlers: list[str] = Field(default_factory=lambda: ["transcribe"])
    available_models: list[str] = Field(default_factory=list)
    model_capabilities: list[dict] = Field(default_factory=list)


class WorkerRegisterResponse(BaseModel):
    worker_id: str
    worker_token: str
    message: str


class ShardStatusResponse(BaseModel):
    shard_id: str
    shard_index: int
    start_seconds: float
    end_seconds: float
    status: str
    claimed_by: str | None = None
    attempts: int = 0


class PollResponse(BaseModel):
    job_id: str
    job_type: str
    audio_url: str
    input_config: dict = Field(default_factory=dict)
    timeout_seconds: int = 3600
    shard_id: str | None = None
    start_seconds: float | None = None
    end_seconds: float | None = None


class BenchmarkResultRequest(BaseModel):
    model: str
    backend: str
    language: str
    audio_seconds: float
    elapsed_seconds: float | None = None
    realtime_factor: float | None = None
    memory_rss_mb: float | None = None
    time_to_first_token_seconds: float | None = None
    word_count: int | None = None
    transcript_preview: str | None = None
    error: str | None = None


class BenchmarkReportRequest(BaseModel):
    created_at: str
    audio_path: str
    machine: dict = Field(default_factory=dict)
    environment: dict = Field(default_factory=dict)
    results: list[BenchmarkResultRequest] = Field(default_factory=list)


class ResultSubmitRequest(BaseModel):
    transcript: str
    language: str = "en"
    model: str = "whisper-large-v3"
    backend: str = "mlx-whisper"
    duration_seconds: float | None = None
    word_count: int | None = None


class FailRequest(BaseModel):
    error: str
