"""每个 ``agentId`` 绑定一个串行工作线程，保证消息按到达顺序处理（对齐 Java ``AgentThreadMgr``）。

Python 版用一组单线程 ``ThreadPoolExecutor`` 来实现：
* 启动时按需创建命名 worker（``agent-worker-N``）；
* ``agent_id`` 与 worker 名字绑定 → 同一 agentId 的所有任务都跑在同一个线程上；
* ``unbind`` 后该 worker 回到 idle 池，可被新 agentId 复用。
"""

from __future__ import annotations

import threading
from concurrent.futures import ThreadPoolExecutor
from typing import Callable, Dict, List, Optional


_WORKER_NAME_PREFIX = "agent-worker-"
_DEFAULT_MAX_WORKERS = 50


class AgentThreadMgr:
    _lock = threading.Lock()
    _name_to_executor: Dict[str, ThreadPoolExecutor] = {}
    _idle_names: List[str] = []
    _agent_id_to_name: Dict[str, str] = {}
    _max_workers: int = _DEFAULT_MAX_WORKERS

    @classmethod
    def configure(cls, max_workers: int) -> None:
        with cls._lock:
            cls._max_workers = max_workers

    @classmethod
    def borrow_idle_thread(cls) -> Optional[str]:
        with cls._lock:
            if cls._idle_names:
                return cls._idle_names.pop()
            if len(cls._name_to_executor) >= cls._max_workers:
                return None
            name = f"{_WORKER_NAME_PREFIX}{len(cls._name_to_executor) + 1}"
            cls._name_to_executor[name] = ThreadPoolExecutor(
                max_workers=1, thread_name_prefix=name,
            )
            return name

    @classmethod
    def bind(cls, agent_id: str, thread_name: str) -> None:
        with cls._lock:
            cls._agent_id_to_name[agent_id] = thread_name

    @classmethod
    def get_thread_name(cls, agent_id: str) -> Optional[str]:
        with cls._lock:
            return cls._agent_id_to_name.get(agent_id)

    @classmethod
    def unbind(cls, agent_id: str) -> Optional[str]:
        with cls._lock:
            name = cls._agent_id_to_name.pop(agent_id, None)
            if name is not None:
                cls._idle_names.append(name)
            return name

    @classmethod
    def submit(cls, thread_name: str, task: Callable[[], None]) -> bool:
        with cls._lock:
            executor = cls._name_to_executor.get(thread_name)
        if executor is None:
            return False
        executor.submit(task)
        return True

    @classmethod
    def shutdown_all(cls) -> None:
        with cls._lock:
            executors = list(cls._name_to_executor.values())
            cls._name_to_executor.clear()
            cls._idle_names.clear()
            cls._agent_id_to_name.clear()
        for ex in executors:
            ex.shutdown(wait=False, cancel_futures=True)

    @classmethod
    def _reset_for_test(cls) -> None:
        cls.shutdown_all()
        cls._max_workers = _DEFAULT_MAX_WORKERS
