state-hub/api/schemas/task.py

import uuid
from datetime import date, datetime
from typing import Self

from pydantic import BaseModel, ConfigDict, field_validator, model_validator

from api.models.task import TaskPriority, TaskStatus
from api.schemas.compat import WorkplanIdCompatMixin, WorkplanIdCreateMixin
from api.task_status import normalize_task_status
from api.work_record_flavor import WORK_RECORD_FLAVORS, normalize_flavor


class TaskStatusMixin(BaseModel):
    @field_validator("status", mode="before", check_fields=False)
    @classmethod
    def _normalize_status(cls, value):
        if value is None:
            return value
        return normalize_task_status(value)


class TaskFlavorMixin(BaseModel):
    @field_validator("flavor", mode="before", check_fields=False)
    @classmethod
    def _normalize_flavor(cls, value):
        flavor = normalize_flavor(value)
        if flavor is not None and flavor not in WORK_RECORD_FLAVORS:
            raise ValueError(
                f"Unknown work-record flavor {flavor!r}; "
                f"expected one of {', '.join(WORK_RECORD_FLAVORS)}"
            )
        return flavor


class TaskCreate(TaskStatusMixin, TaskFlavorMixin, WorkplanIdCreateMixin):
    id: uuid.UUID | None = None
    record_id: str | None = None
    flavor: str | None = None
    title: str
    description: str | None = None
    status: TaskStatus = TaskStatus.todo
    priority: TaskPriority = TaskPriority.medium
    assignee: str | None = None
    due_date: date | None = None
    blocking_reason: str | None = None
    needs_human: bool = False
    intervention_note: str | None = None
    parent_task_id: uuid.UUID | None = None

    @model_validator(mode="after")
    def intervention_note_required_when_flagged(self) -> Self:
        if self.needs_human and not self.intervention_note:
            raise ValueError("intervention_note is required when needs_human is True")
        return self


class TaskUpdate(TaskStatusMixin, TaskFlavorMixin):
    title: str | None = None
    flavor: str | None = None
    description: str | None = None
    status: TaskStatus | None = None
    priority: TaskPriority | None = None
    assignee: str | None = None
    due_date: date | None = None
    blocking_reason: str | None = None
    needs_human: bool | None = None
    intervention_note: str | None = None
    parent_task_id: uuid.UUID | None = None
    # Token passthrough — three tiers (highest precision wins):
    # 1. tokens_in + tokens_out  → exact counts; note defaults to "measured"
    # 2. workplan_tokens_in + workplan_tokens_out  → prorated across task count (note="workplan")
    # 3. neither provided, status=done  → heuristic 1000/500 (note="heuristic")
    # token_note overrides the auto-assigned note for Tier 1 only (e.g. "userbased")
    # suppress_token_event lets file/cache sync update status without recording usage.
    tokens_in: int | None = None
    tokens_out: int | None = None
    workplan_tokens_in: int | None = None
    workplan_tokens_out: int | None = None
    token_note: str | None = None
    model: str | None = None
    agent: str | None = None
    session_id: str | None = None
    suppress_token_event: bool | None = None

    @model_validator(mode="after")
    def blocking_reason_required_when_human_waiting(self) -> Self:
        if self.status == TaskStatus.wait and self.needs_human and not self.blocking_reason:
            raise ValueError("blocking_reason is required when a human-blocked task is waiting")
        return self

    @model_validator(mode="after")
    def intervention_note_required_when_flagged(self) -> Self:
        if self.needs_human and not self.intervention_note:
            raise ValueError("intervention_note is required when needs_human is True")
        return self


class TaskStatusBulkUpdate(TaskStatusMixin):
    task_id: uuid.UUID
    status: TaskStatus
    blocking_reason: str | None = None


class TaskStatusBulkSync(BaseModel):
    updates: list[TaskStatusBulkUpdate]
    author: str | None = "custodian"
    session_id: str | None = None

    @field_validator("updates")
    @classmethod
    def updates_required(cls, value: list[TaskStatusBulkUpdate]):
        if not value:
            raise ValueError("at least one task status update is required")
        return value


class TaskRead(TaskStatusMixin, WorkplanIdCompatMixin):
    model_config = ConfigDict(from_attributes=True)
    id: uuid.UUID
    record_id: str | None = None
    flavor: str | None = None
    title: str
    description: str | None = None
    status: TaskStatus
    priority: TaskPriority
    assignee: str | None = None
    due_date: date | None = None
    blocking_reason: str | None = None
    needs_human: bool
    intervention_note: str | None = None
    parent_task_id: uuid.UUID | None = None
    created_at: datetime
    updated_at: datetime


class TaskCountRead(TaskStatusMixin, WorkplanIdCompatMixin):
    status: TaskStatus
    count: int


class TaskStatusBulkSyncRead(BaseModel):
    updated: list[TaskRead]
    progress_event_ids: list[uuid.UUID]