# session_manager.py（最终版）
import asyncio
import logging
from dataclasses import dataclass, field
from fastapi import WebSocket
from typing import Optional, List, Any
from asyncio import Queue

logger = logging.getLogger(__name__)

@dataclass
class WsSession:
    """WebSocket 会话类：管理单个连接的所有上下文（无用户标识，仅做资源隔离）"""
    websocket: WebSocket                  # 前端 WS 连接
    session_id: str                       # 会话ID（用 conn_id 即可，无需用户ID）
    is_closed: bool = False               # 会话是否关闭
    tasks: List[asyncio.Task] = field(default_factory=list)  # 转发协程列表
    # ASR/TTS 客户端连接（每个会话独立创建，天然隔离）
    asr_client: Optional[Any] = None      # ASR 客户端任务（存储协程任务）
    tts_client: Optional[Any] = None      # TTS 客户端实例（无用户态）
    # 新增：ASR 依赖的核心属性
    audio_queue: Queue = field(default_factory=lambda: Queue(maxsize=30))  # 音频队列（防积压）
    llm_queue: Queue = field(default_factory=Queue)                        # LLM 队列（转发ASR最终结果）
    asr_ws_task: Optional[asyncio.Task] = None                            # ASR 核心协程任务

async def create_session(websocket: WebSocket) -> WsSession:
    """创建会话（仅初始化资源，无用户相关逻辑）"""
    conn_id = str(id(websocket))
    session = WsSession(
        websocket=websocket,
        session_id=conn_id  # 会话ID = 连接ID，无需用户标识
    )
    logger.info(f"会话 [{session.session_id}] 创建成功")
    return session

async def close_session(session: WsSession):
    """关闭会话：清理所有协程和ASR/TTS连接（核心：资源释放）"""
    session.is_closed = True
    # 1. 取消所有转发协程
    for task in session.tasks:
        if not task.done():
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
    # 2. 单独取消ASR核心任务（新增）
    if session.asr_ws_task and not session.asr_ws_task.done():
        session.asr_ws_task.cancel()
        try:
            await session.asr_ws_task
        except asyncio.CancelledError:
            logger.info(f"会话 [{session.session_id}] ASR核心协程已取消")
    # 3. 关闭ASR/TTS客户端（适配：asr_client 是任务，无需close，仅日志提示）
    if session.asr_client:
        logger.info(f"会话 [{session.session_id}] ASR 客户端任务已清理")
    if session.tts_client and hasattr(session.tts_client, 'closed') and not session.tts_client.closed:
        await session.tts_client.close()
        logger.info(f"会话 [{session.session_id}] TTS 客户端已关闭")
    # 4. 清空音频队列（防止内存泄漏）
    while not session.audio_queue.empty():
        try:
            session.audio_queue.get_nowait()
            session.audio_queue.task_done()
        except asyncio.QueueEmpty:
            break
    logger.info(f"会话 [{session.session_id}] 已完全清理")