"""In-memory state backend for task persistence.

Useful for testing and development when no database is needed.
"""

from __future__ import annotations

from typing import Any

from fusionagi.core.persistence import StateBackend
from fusionagi.schemas.task import Task, TaskState


class InMemoryStateBackend(StateBackend):
    """In-memory implementation of StateBackend.

    All data is lost on process restart. Use SQLiteStateBackend
    or a Postgres-backed backend for production persistence.
    """

    def __init__(self) -> None:
        self._tasks: dict[str, Task] = {}
        self._traces: dict[str, list[dict[str, Any]]] = {}

    def get_task(self, task_id: str) -> Task | None:
        """Load task by id."""
        return self._tasks.get(task_id)

    def set_task(self, task: Task) -> None:
        """Save task."""
        self._tasks[task.task_id] = task

    def get_task_state(self, task_id: str) -> TaskState | None:
        """Return current task state or None if task unknown."""
        task = self._tasks.get(task_id)
        return task.state if task else None

    def set_task_state(self, task_id: str, state: TaskState) -> None:
        """Update task state; creates no task if missing."""
        task = self._tasks.get(task_id)
        if task is not None:
            self._tasks[task_id] = task.model_copy(update={"state": state})

    def append_trace(self, task_id: str, entry: dict[str, Any]) -> None:
        """Append trace entry."""
        if task_id not in self._traces:
            self._traces[task_id] = []
        self._traces[task_id].append(entry)

    def get_trace(self, task_id: str) -> list[dict[str, Any]]:
        """Load trace for task."""
        return list(self._traces.get(task_id, []))

    def list_tasks(self, state: TaskState | None = None, limit: int = 100) -> list[Task]:
        """List tasks, optionally filtered by state."""
        tasks = list(self._tasks.values())
        if state is not None:
            tasks = [t for t in tasks if t.state == state]
        return tasks[:limit]

    def delete_task(self, task_id: str) -> bool:
        """Delete a task and its traces."""
        self._traces.pop(task_id, None)
        return self._tasks.pop(task_id, None) is not None

    def count_tasks(self) -> int:
        """Return total task count."""
        return len(self._tasks)