from pydantic import BaseModel, field_validator, model_validator


class SmartTestState(BaseModel):
    state: str = "idle"
    percent: int | None = None
    eta_seconds: int | None = None
    eta_timestamp: str | None = None
    started_at: str | None = None
    finished_at: str | None = None
    error_text: str | None = None


_VALID_STAGE_NAMES = frozenset({"surface_validate", "short_smart", "long_smart"})


class StartBurninRequest(BaseModel):
    drive_ids: list[int]
    operator: str
    run_surface: bool = True
    run_short:   bool = True
    run_long:    bool = True
    stage_order: list[str] | None = None  # custom execution order, e.g. ["short_smart","long_smart","surface_validate"]

    @field_validator("operator")
    @classmethod
    def validate_operator(cls, v: str) -> str:
        v = v.strip()
        if not v:
            raise ValueError("operator must not be empty")
        return v

    @model_validator(mode="after")
    def validate_stages(self) -> "StartBurninRequest":
        if not (self.run_surface or self.run_short or self.run_long):
            raise ValueError("At least one stage must be selected")
        if self.stage_order is not None:
            for s in self.stage_order:
                if s not in _VALID_STAGE_NAMES:
                    raise ValueError(f"Invalid stage name in stage_order: {s!r}")
        return self

    @property
    def profile(self) -> str:
        _MAP = {
            (True,  True,  True):  "full",
            (True,  True,  False): "surface_short",
            (True,  False, True):  "surface_long",
            (True,  False, False): "surface",
            (False, True,  True):  "short_long",
            (False, True,  False): "short",
            (False, False, True):  "long",
        }
        return _MAP[(self.run_surface, self.run_short, self.run_long)]


class CancelBurninRequest(BaseModel):
    operator: str = "unknown"


class BurninStageResponse(BaseModel):
    id: int
    stage_name: str
    state: str
    percent: int = 0
    started_at: str | None = None
    finished_at: str | None = None
    error_text: str | None = None


class BurninJobResponse(BaseModel):
    id: int
    drive_id: int
    profile: str
    state: str
    percent: int = 0
    stage_name: str | None = None
    operator: str
    created_at: str
    started_at: str | None = None
    finished_at: str | None = None
    error_text: str | None = None
    stages: list[BurninStageResponse] = []


class DriveResponse(BaseModel):
    id: int
    devname: str
    serial: str | None = None
    model: str | None = None
    size_bytes: int | None = None
    temperature_c: int | None = None
    smart_health: str = "UNKNOWN"
    last_polled_at: str
    is_stale: bool
    smart_short: SmartTestState
    smart_long: SmartTestState
    notes: str | None = None
    location: str | None = None


class UpdateDriveRequest(BaseModel):
    notes: str | None = None
    location: str | None = None