import asyncio
from services.audio_session import AudioSession
from services.task_manager import task_manager
from utils.logger import session_logger
from services.interrupt_handler import interrupt_handler
import time
from config import settings


async def _connect_with_interrupt_cleanup(service, session_key):
    """Own warmup connections even when interruption races with connect()."""
    if interrupt_handler.is_interrupted(session_key):
        return False
    await interrupt_handler.register_cleanup_handler(
        session_key, service.close,
        priority=interrupt_handler.PRIORITY["NETWORK_CONNECTION"],
    )
    try:
        await service.connect()
        if interrupt_handler.is_interrupted(session_key):
            await service.close()
            return False
        return True
    except asyncio.CancelledError:
        await service.close()
        raise
    except Exception:
        await service.close()
        raise


async def prepare_asr_service(session, device_id, session_id):
    session.asr_error = None
    session.asr_service.sample_rate = session.sample_rate
    if hasattr(session.asr_service, "device_id"):
        session.asr_service.device_id = device_id
    if hasattr(session.asr_service, "session_id"):
        session.asr_service.session_id = session_id
    await establish_asr_connection(session)


async def prepare_llm_service(session, device_id, session_id, selected_role):
    if hasattr(session.llm_service, "device_id"):
        session.llm_service.device_id = device_id
    if hasattr(session.llm_service, "session_id"):
        session.llm_service.session_id = session_id
    if hasattr(session.llm_service, "connect"):
        if not await _connect_with_interrupt_cleanup(session.llm_service, (device_id, session_id)):
            return
    
    from services.conversation_history import conversation_history_manager
    from config import settings
    device_history = await conversation_history_manager.get_history(device_id)
    if device_history:
        session.llm_service.history = device_history.history[
            -settings.max_conversation_history :
        ]


async def prepare_tts_service(session, device_id, session_id, selected_role):
    session.tts_service.selected_role = selected_role
    if hasattr(session.tts_service, "device_id"):
        session.tts_service.device_id = device_id
    if hasattr(session.tts_service, "session_id"):
        session.tts_service.session_id = session_id
    await establish_tts_connection(session)


async def establish_asr_connection(session: AudioSession):

    try:
        if session.device_id and session.session_id:
            session_key = (session.device_id, session.session_id)
            if interrupt_handler.is_interrupted(session_key):
                session_logger.info(
                    session.device_id, session.session_id, "会话已中断，跳过建立ASR连接"
                )
                return
        session.send_task = await task_manager.create_task(
            send_audio_task(session),
            device_id=session.device_id,
            session_key=(session.device_id, session.session_id) if session.device_id and session.session_id else None,
            task_type="audio_send"
        )
        await session.asr_service.connect()
        if not session.asr_service.connected:
            raise RuntimeError("语音识别 WebSocket 连接未建立")
        if session.device_id and session.session_id:
            if interrupt_handler.is_interrupted((session.device_id, session.session_id)):
                raise RuntimeError("连接期间会话已中断")
            await interrupt_handler.register_cleanup_handler(
                (session.device_id, session.session_id), session.asr_service.close,
                priority=interrupt_handler.PRIORITY["NETWORK_CONNECTION"],
            )
            await interrupt_handler.register_cleanup_handler(
                (session.device_id, session.session_id), session.cancel_send_task,
                priority=interrupt_handler.PRIORITY["TASK_CANCELLATION"],
            )
        session_logger.info(
            "unknown", "unknown", f"语音识别连接已建立，采样率: {session.sample_rate}Hz"
        )
    except asyncio.CancelledError:
        await session.cancel_send_task()
        await session.asr_service.close()
        raise
    except Exception as e:
        session.asr_error = f"{type(e).__name__}: {e}"
        await session.cancel_send_task()
        await session.asr_service.close()
        session_logger.error(
            "unknown", "unknown", f"建立语音识别 WebSocket 连接失败: {e}", exc_info=True
        )


async def establish_tts_connection(session: AudioSession):

    try:
        if session.device_id and session.session_id:
            session_key = (session.device_id, session.session_id)
            if interrupt_handler.is_interrupted(session_key):
                session_logger.info(
                    session.device_id, session.session_id, "会话已中断，跳过建立TTS连接"
                )
                return

        if not await _connect_with_interrupt_cleanup(
            session.tts_service, (session.device_id, session.session_id)
        ):
            return

        if not session.tts_service.connected:
            session_logger.error(
                session.device_id or "unknown",
                session.session_id or "unknown",
                "语音合成连接未建立",
            )
            return

        session_logger.info(
            session.device_id or "unknown",
            session.session_id or "unknown",
            "语音合成连接已建立",
        )
    except Exception as e:
        session_logger.error(
            session.device_id or "unknown",
            session.session_id or "unknown",
            f"建立语音合成连接失败: {e}",
        )


async def send_audio_task(session: AudioSession):
    try:
        while True:
            if session.device_id and session.session_id:
                session_key = (session.device_id, session.session_id)
                if interrupt_handler.is_interrupted(session_key):
                    session_logger.info(
                        session.device_id, session.session_id, "发送音频任务被中断"
                    )
                    break
            queue = session.send_queue
            audio_data = await queue.get()
            try:
                if audio_data is None:  # Drain all preceding audio before finish-task.
                    break
                await session.asr_service.send_audio(audio_data)
            finally:
                queue.task_done()

    except asyncio.CancelledError:
        session.asr_error = session.asr_error or "发送音频任务被取消"
        session_logger.info(
            session.device_id or "unknown",
            session.session_id or "unknown",
            "发送音频数据任务被取消",
        )
    except Exception as e:
        session.asr_error = str(e)
        session_logger.error(
            session.device_id or "unknown",
            session.session_id or "unknown",
            f"发送音频数据时出错: {str(e)}",
            exc_info=True,
        )


async def end_asr_session(session: AudioSession):
    """Drain device audio, collect the complete transcript, always release ASR."""
    session.transcript = ""
    try:
        session_key = (session.device_id, session.session_id)
        if interrupt_handler.is_interrupted(session_key):
            raise RuntimeError("识别会话已中断")
        if session.send_task and not session.send_task.done():
            await session.send_queue.put(None)
            await asyncio.wait_for(
                asyncio.shield(session.send_task),
                timeout=settings.asr_queue_drain_timeout_seconds,
            )
        if session.asr_error:
            raise RuntimeError(session.asr_error)
        if session.send_queue and not session.send_queue.empty():
            raise RuntimeError("发送任务提前退出，音频未完整发送")
        if interrupt_handler.is_interrupted(session_key):
            raise RuntimeError("识别会话已中断")
        asr_end_start_time = time.perf_counter()
        await session.asr_service.send_end()
        if interrupt_handler.is_interrupted(session_key):
            raise RuntimeError("识别会话已中断")
        session.transcript = await session.asr_service.receive_results()
        session_logger.info(
            session.device_id or "unknown", session.session_id or "unknown",
            f"ASR结束延迟: {time.perf_counter() - asr_end_start_time:.2f} 秒",
        )
    except Exception as e:
        session.asr_error = f"{type(e).__name__}: {e}"
        session.transcript = ""
        session_logger.error(
            session.device_id or "unknown", session.session_id or "unknown",
            f"ASR结束处理出错: {session.asr_error}",
        )
    finally:
        if session.send_task:
            session.send_task.cancel()
            try:
                await asyncio.wait_for(session.send_task, timeout=2.0)
            except (asyncio.CancelledError, asyncio.TimeoutError):
                pass
            finally:
                session.send_task = None
        await session.asr_service.close()
        session_logger.info(
            session.device_id or "unknown", session.session_id or "unknown",
            "语音识别 WebSocket 连接和 ClientSession 已关闭",
        )