import asyncio
import logging
import time
from typing import Dict, Optional, Any, List, Callable, Awaitable
from dataclasses import dataclass, field

logger = logging.getLogger("telepush.worker")

# Глобальные настройки лимитов
MAX_CONCURRENT_PER_USER = 15  # Задач на одного юзера одновременно
MAX_JOIN_CONCURRENT_PER_USER = 15  # Лимит специально для вступлений (join)
_GLOBAL_SEMAPHORE_SIZE = 300  # Максимум задач на весь бот сразу

# Реестры состояний
_user_pools: Dict[int, "UserWorkerPool"] = {}
_global_semaphore: Optional[asyncio.Semaphore] = None
_user_join_semaphores: Dict[int, asyncio.Semaphore] = {}


def _get_global_semaphore() -> asyncio.Semaphore:
    global _global_semaphore
    if _global_semaphore is None:
        _global_semaphore = asyncio.Semaphore(_GLOBAL_SEMAPHORE_SIZE)
    return _global_semaphore


def get_join_semaphore(user_id: int) -> asyncio.Semaphore:
    if user_id not in _user_join_semaphores:
        _user_join_semaphores[user_id] = asyncio.Semaphore(MAX_JOIN_CONCURRENT_PER_USER)
    return _user_join_semaphores[user_id]


@dataclass
class TaskInfo:
    task_id: str
    task: asyncio.Task
    job_id: Optional[str]
    created_at: float = field(default_factory=time.time)
    user_id: int = 0


class UserWorkerPool:
    """Пул задач с изоляцией ресурсов по пользователям."""

    def __init__(self, user_id: int, max_concurrent: int = MAX_CONCURRENT_PER_USER):
        self.user_id = user_id
        self._semaphore = asyncio.Semaphore(max_concurrent)
        self._tasks: Dict[str, TaskInfo] = {}
        self._task_counter = 0

    def submit(
        self,
        coro_factory: Callable[[], Awaitable[Any]],
        job_id: str = None,
        name: str = None,
    ) -> asyncio.Task:
        """Регистрирует корутину и запускает её в асинхронном фоне."""
        self._task_counter += 1
        t_id = f"uid{self.user_id}_{self._task_counter}"

        async def _wrapper():
            global_sem = _get_global_semaphore()
            try:
                # Асинхронно ждем доступа к ресурсам (сначала своим, потом общим)
                async with self._semaphore, global_sem:
                    coro = coro_factory()
                    return await coro
            except asyncio.CancelledError:
                # Пробрасываем отмену для корректного завершения Event Loop
                raise
            except Exception as e:
                logger.error(
                    f"[pool] Ошибка в задаче {job_id or t_id} (uid={self.user_id}): {e}"
                )
            finally:
                # Удаляем задачу из списка активных
                self._tasks.pop(t_id, None)
                self._cleanup_if_idle()

        # Создаем задачу и сохраняем ссылку на неё
        task = asyncio.create_task(_wrapper(), name=name or t_id)
        self._tasks[t_id] = TaskInfo(t_id, task, job_id, user_id=self.user_id)
        return task

    def cancel_all(self):
        """Отменяет все задачи этого пользователя."""
        for info in list(self._tasks.values()):
            if info.task and not info.task.done():
                info.task.cancel()

    def active_count(self) -> int:
        return len(self._tasks)

    def _cleanup_if_idle(self):
        """Удаляет объект пула, если в нем нет активных задач (экономия RAM)."""
        if not self._tasks:
            _user_pools.pop(self.user_id, None)
            _user_join_semaphores.pop(self.user_id, None)


# --- Глобальный API (используется в хэндлерах) ---


def get_user_pool(user_id: int) -> UserWorkerPool:
    if user_id not in _user_pools:
        _user_pools[user_id] = UserWorkerPool(user_id)
    return _user_pools[user_id]


def submit_for_user(
    user_id: int, coro_factory: Callable, job_id: str = None, name: str = None
) -> asyncio.Task:
    return get_user_pool(user_id).submit(coro_factory, job_id, name)


def submit_join_for_user(
    user_id: int, coro_factory: Callable, job_id: str = None, name: str = None
) -> asyncio.Task:
    """Версия для Join-задач с дополнительным ограничением семафора вступлений."""
    join_sem = get_join_semaphore(user_id)

    async def wrapped_factory():
        async with join_sem:
            return await coro_factory()

    return get_user_pool(user_id).submit(wrapped_factory, job_id, name)


def get_all_active_counts() -> Dict[int, int]:
    return {uid: p.active_count() for uid, p in _user_pools.items()}


def total_active_tasks() -> int:
    return sum(p.active_count() for p in _user_pools.values())


async def shutdown_all_pools():
    """Асинхронно останавливает все задачи во всех пулах."""
    all_tasks = []
    for pool in list(_user_pools.values()):
        for info in pool._tasks.values():
            if not info.task.done():
                info.task.cancel()
                all_tasks.append(info.task)

    if all_tasks:
        logger.info(f"[pool] Завершение {len(all_tasks)} фоновых задач...")
        # Даем задачам время "переварить" отмену (CancelledError)
        await asyncio.gather(*all_tasks, return_exceptions=True)
