"""In-memory task registry for background dispatch tasks.

One global instance per agent process. Workers add steps as they progress;
the agent subscribes to push updates to the room (UI) and to speak results.
"""
from __future__ import annotations

import asyncio
import time
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable


class TaskStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class TaskStep:
    role: str  # "thinking", "tool_call", "tool_result", "text"
    content: str
    timestamp: float = field(default_factory=time.time)


@dataclass
class Task:
    id: str
    description: str
    status: TaskStatus = TaskStatus.PENDING
    steps: list[TaskStep] = field(default_factory=list)
    result: str | None = None
    error: str | None = None
    created_at: float = field(default_factory=time.time)
    completed_at: float | None = None


class TaskRegistry:
    """Thread-safe registry of background tasks with pub/sub notifications."""

    def __init__(self):
        self._tasks: dict[str, Task] = {}
        self._listeners: list[Callable[[str, dict], Any]] = []
        self._lock = asyncio.Lock()

    async def create(self, description: str) -> Task:
        task = Task(id=uuid.uuid4().hex[:8], description=description)
        async with self._lock:
            self._tasks[task.id] = task
        await self._notify(task.id, {"event": "created", "task": self._task_dict(task)})
        return task

    async def add_step(self, task_id: str, role: str, content: str):
        async with self._lock:
            task = self._tasks.get(task_id)
            if not task:
                return
            task.steps.append(TaskStep(role=role, content=content))
            if task.status == TaskStatus.PENDING:
                task.status = TaskStatus.RUNNING
        await self._notify(task_id, {"event": "step", "role": role, "content": content})

    async def complete(self, task_id: str, result: str):
        async with self._lock:
            task = self._tasks.get(task_id)
            if not task:
                return
            task.status = TaskStatus.COMPLETED
            task.result = result
            task.completed_at = time.time()
        await self._notify(task_id, {"event": "completed", "result": result})

    async def fail(self, task_id: str, error: str):
        async with self._lock:
            task = self._tasks.get(task_id)
            if not task:
                return
            task.status = TaskStatus.FAILED
            task.error = error
            task.completed_at = time.time()
        await self._notify(task_id, {"event": "failed", "error": error})

    async def get(self, task_id: str) -> dict | None:
        async with self._lock:
            t = self._tasks.get(task_id)
            return self._task_dict(t) if t else None

    async def list_all(self) -> list[dict]:
        async with self._lock:
            return [self._task_dict(t) for t in self._tasks.values()]

    def subscribe(self, callback: Callable[[str, dict], Any]):
        self._listeners.append(callback)

    async def _notify(self, task_id: str, event: dict):
        for cb in self._listeners:
            try:
                result = cb(task_id, event)
                if asyncio.iscoroutine(result):
                    await result
            except Exception:
                pass

    @staticmethod
    def _task_dict(task: Task) -> dict:
        return {
            "id": task.id,
            "description": task.description,
            "status": task.status.value,
            "steps": [
                {"role": s.role, "content": s.content, "ts": s.timestamp}
                for s in task.steps
            ],
            "result": task.result,
            "error": task.error,
            "created_at": task.created_at,
            "completed_at": task.completed_at,
        }


# Global singleton — one per agent process
registry = TaskRegistry()