from dataclasses import dataclass, field
from threading import Lock
from uuid import uuid4

from app.schemas.changeset import ChangeItem
from app.schemas.chat import TaskResultType, TaskStatus
from app.schemas.common import ErrorPayload


@dataclass
class TaskState:
    task_id: str
    status: TaskStatus = TaskStatus.QUEUED
    result_type: TaskResultType | None = None
    answer: str | None = None
    changeset: list[ChangeItem] = field(default_factory=list)
    error: ErrorPayload | None = None


class TaskStore:
    def __init__(self) -> None:
        self._items: dict[str, TaskState] = {}
        self._lock = Lock()

    def create(self) -> TaskState:
        task = TaskState(task_id=str(uuid4()))
        with self._lock:
            self._items[task.task_id] = task
        return task

    def get(self, task_id: str) -> TaskState | None:
        with self._lock:
            return self._items.get(task_id)

    def save(self, task: TaskState) -> None:
        with self._lock:
            self._items[task.task_id] = task