from __future__ import annotations

from datetime import datetime
from enum import StrEnum
from typing import Any
from uuid import uuid4

from pydantic import BaseModel, ConfigDict, Field


class SampleStatus(StrEnum):
    active = "active"
    superseded = "superseded"


class PendingStatus(StrEnum):
    pending = "pending"
    leased = "leased"
    promoted = "promoted"
    failed = "failed"


class PoolSample(BaseModel):
    """A single pool sample row."""

    model_config = ConfigDict(frozen=True)

    sample_id: str = Field(default_factory=lambda: uuid4().hex)
    sample_idx: int | None = None
    key_values: dict[str, Any] = Field(default_factory=dict)
    payload: dict[str, Any] = Field(default_factory=dict)
    source_run_id: str | None = None
    call_id: str | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)
    status: SampleStatus = SampleStatus.active
    created_at: datetime | None = None


class AcquireQuery(BaseModel):
    """Query for no-replacement sample acquisition."""

    model_config = ConfigDict(frozen=True)

    run_id: str
    request_id: str = Field(default_factory=lambda: uuid4().hex)
    key_values: dict[str, Any] = Field(default_factory=dict)
    n: int
    consumer_tag: str = ""


class AcquireResult(BaseModel):
    """Result of a sample acquisition."""

    model_config = ConfigDict(frozen=True)

    samples: list[PoolSample] = Field(default_factory=list)
    claimed: int = 0

    def deficit(self, requested_n: int) -> int:
        return max(requested_n - len(self.samples), 0)


class PoolClaim(BaseModel):
    """Claim record for no-replacement tracking."""

    model_config = ConfigDict(frozen=True)

    claim_id: str = Field(default_factory=lambda: uuid4().hex)
    run_id: str
    request_id: str
    consumer_tag: str = ""
    sample_id: str
    claim_idx: int
    claimed_at: datetime | None = None


class PendingSample(BaseModel):
    """Sample in pending state awaiting validation/promotion."""

    model_config = ConfigDict(frozen=True)

    pending_id: str = Field(default_factory=lambda: uuid4().hex)
    key_values: dict[str, Any] = Field(default_factory=dict)
    sample_idx: int = 0
    payload: dict[str, Any] = Field(default_factory=dict)
    source_run_id: str | None = None
    call_id: str | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)
    priority: int = 0
    status: PendingStatus = PendingStatus.pending
    worker_id: str | None = None
    lease_expires_at: datetime | None = None
    attempt_count: int = 0
    created_at: datetime | None = None


class InsertResult(BaseModel):
    """Result of a bulk insert operation."""

    model_config = ConfigDict(frozen=True)

    inserted: int = 0
    skipped: int = 0
    failed: int = 0


class PendingStatusCounts(BaseModel):
    """Counts of pending rows by lifecycle status."""

    model_config = ConfigDict(frozen=True)

    pending: int = 0
    leased: int = 0
    promoted: int = 0
    failed: int = 0

    @property
    def total(self) -> int:
        return self.pending + self.leased + self.promoted + self.failed


class WorkerSnapshot(BaseModel):
    """Observable state for a running pool worker controller."""

    model_config = ConfigDict(frozen=True)

    worker_count: int
    stop_requested: bool = False
    claimed: int = 0
    promoted: int = 0
    failed: int = 0
    retried: int = 0
    process_errors: int = 0
    idle_polls: int = 0
    status_counts: PendingStatusCounts = Field(default_factory=PendingStatusCounts)


class CoverageRow(BaseModel):
    """Aggregate count per unique key combination."""

    model_config = ConfigDict(frozen=True)

    key_values: dict[str, Any] = Field(default_factory=dict)
    count: int = 0
