"""Shared ownership of peer connections and their pipeline tasks."""

from __future__ import annotations

import asyncio
from collections.abc import Coroutine
from typing import Any

from loguru import logger


active_connections: set[object] = set()
pipeline_tasks: set[asyncio.Task[None]] = set()
connection_tasks: dict[object, asyncio.Task[None]] = {}
DEFAULT_CLOSE_GRACE_SECONDS = 10.0


def _consume_pipeline_result(task: asyncio.Task[None], connection: object) -> None:
    pipeline_tasks.discard(task)
    if connection_tasks.get(connection) is task:
        connection_tasks.pop(connection, None)
    try:
        error = task.exception()
    except asyncio.CancelledError:
        logger.info(f"Realtime pipeline 已取消: task={task.get_name()}")
        return
    if error is not None:
        logger.opt(exception=error).error(
            f"Realtime pipeline 异常结束: task={task.get_name()}"
        )


def start_pipeline_task(
    connection: object,
    coroutine: Coroutine[Any, Any, None],
    *,
    protocol: str,
) -> asyncio.Task[None]:
    connection_id = str(getattr(connection, "pc_id", id(connection)))
    task = asyncio.create_task(
        coroutine,
        name=f"{protocol}-pipeline:{connection_id}",
    )
    active_connections.add(connection)
    pipeline_tasks.add(task)
    connection_tasks[connection] = task
    task.add_done_callback(
        lambda completed, connection=connection: _consume_pipeline_result(
            completed,
            connection,
        )
    )
    return task


async def wait_for_pipeline_close(
    task: asyncio.Task[None] | None,
    *,
    connection_id: str,
    timeout: float = DEFAULT_CLOSE_GRACE_SECONDS,
) -> None:
    if task is None:
        return
    try:
        await asyncio.wait_for(asyncio.shield(task), timeout=timeout)
        return
    except TimeoutError:
        logger.warning(
            f"Realtime pipeline 关闭超过 {timeout:g} 秒，执行取消: "
            f"connection_id={connection_id}"
        )
    except asyncio.CancelledError:
        raise
    except Exception:
        return

    task.cancel()
    done, _pending = await asyncio.wait({task}, timeout=timeout)
    if not done:
        logger.error(f"Realtime pipeline 取消后仍未退出: connection_id={connection_id}")


async def shutdown_active_sessions(
    *,
    timeout: float = DEFAULT_CLOSE_GRACE_SECONDS,
) -> None:
    connections = list(active_connections)
    if connections:
        await asyncio.gather(
            *(connection.disconnect() for connection in connections),
            return_exceptions=True,
        )
        active_connections.difference_update(connections)

    tasks = list(pipeline_tasks)
    if not tasks:
        return
    done, pending = await asyncio.wait(tasks, timeout=timeout)
    for task in pending:
        task.cancel()
    cancelled: set[asyncio.Task[None]] = set()
    if pending:
        cancelled, stuck = await asyncio.wait(pending, timeout=timeout)
        if stuck:
            logger.error(f"应用关闭时仍有 {len(stuck)} 个 Realtime pipeline 未退出")
    logger.info(
        f"Realtime 会话清理完成: normal={len(done)} cancelled={len(cancelled)}"
    )