from __future__ import annotations

from collections.abc import AsyncIterator
from typing import Any, Mapping, Optional

from app.core.task_manager.stores import TaskStore
from app.core.task_manager.strategies import DeliveryStrategy
from app.core.task_manager.types import BaseTask, DeliveryMode, TaskRecord, TaskStatusView


class TaskManager:
    """任务管理器：对上层（路由/业务）提供统一入口。

    init params:
    - store: TaskStore, 任务存储

    特性：
    - store 可插拔：内存 or MySQL（通过 SQLAlchemy）
    - delivery strategy 可插拔：streaming / async_polling / ...（后续扩展只需新增策略并注册）
    - 为高频轮询提供轻量查询：get_status()
    """

    def __init__(self, *, store: TaskStore, strategies: Mapping[DeliveryMode, DeliveryStrategy]) -> None:
        self.store = store
        self._strategies = dict(strategies)

    def _strategy_for(self, mode: DeliveryMode) -> DeliveryStrategy:
        strategy = self._strategies.get(mode)
        if strategy is None:
            raise ValueError(f"No delivery strategy registered for mode={mode!r}")
        return strategy

    async def create(
        self,
        *,
        task: BaseTask,
        mode: DeliveryMode,
        run_args: dict[str, Any] | None = None,
    ) -> TaskRecord:
        """创建任务记录。

        参数：
        - task: 实际要执行的 Task 对象，必须实现 run/status/is_done/get_result
        - mode: 结果交付方式（流式 / 任务+轮询）
        - run_args: 传给 task.run 的参数（会序列化到 payload，便于后续 worker 使用）
        """
        payload: dict[str, Any] = {
            "task_class": task.__class__.__name__,
            "run_args": run_args or {},
        }
        return await self.store.create(payload=payload, mode=mode)

    async def start(self, *, task_id: str) -> Optional[AsyncIterator[Any]]:
        task = await self.store.get(task_id)
        if task is None:
            raise ValueError(f"Task not found: {task_id}")
        strategy = self._strategy_for(task.mode)
        return await strategy.start(task)

    async def get_status(self, *, task_id: str) -> TaskStatusView:
        view = await self.store.get_status_view(task_id)
        if view is None:
            raise ValueError(f"Task not found: {task_id}")
        return view

    async def stream(self, *, task_id: str) -> AsyncIterator[Any]:
        task = await self.store.get(task_id)
        if task is None:
            raise ValueError(f"Task not found: {task_id}")
        strategy = self._strategy_for(task.mode)
        return await strategy.stream(task)
