import asyncio
import json

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from app.core.database import async_db_session
from app.core.dependencies import _authenticate
from app.core.exceptions import CustomException
from app.core.logger import logger
from app.core.router_class import OperationLogRoute

from .schema import ChatQuerySchema
from .service import ChatService, get_user_model_config

WS_AI = APIRouter(
    route_class=OperationLogRoute,
    prefix="/ai/chat",
    tags=["智能助手WebSocket"],
)


async def _send_error_and_close(websocket: WebSocket, message: str) -> None:
    """发送错误消息并关闭连接"""
    try:
        await websocket.send_text(f"错误: {message}")
    except RuntimeError:
        pass
    finally:
        try:
            await websocket.close()
        except RuntimeError:
            pass


@WS_AI.websocket("/ws", name="WebSocket聊天")
async def websocket_chat_controller(websocket: WebSocket) -> None:
    """
    WebSocket 聊天接口。

    支持的消息格式（JSON）：
    - 对话：{"message": "...", "session_id": "...", "files": [...]}
    - 停止：{"action": "stop", "session_id": "..."}

    ws://127.0.0.1:8001/api/v1/ai/chat/ws?token=xxx
    """
    await websocket.accept()
    token = websocket.query_params.get("token")

    if not token:
        await _send_error_and_close(websocket, "未提供认证token，请重新登录")
        return

    # 跨消息循环共享的停止信号：客户端发送 stop 时 set，生成器检测到后退出
    stop_event = asyncio.Event()
    # 标记当前是否在生成中，便于 stop 校验
    is_generating = asyncio.Event()

    try:
        redis = websocket.app.state.redis
        async with async_db_session() as db:
            auth = await _authenticate(token, db, redis)

            user = auth.user
            logger.info("WebSocket连接已建立: {} - 用户: {}", websocket.client, user.username if user else "未认证")

            chat_service = ChatService(auth)

            # 消息循环
            while True:
                try:
                    data = await websocket.receive_text()
                    try:
                        message_data = json.loads(data)
                        query = ChatQuerySchema(**message_data)
                    except json.JSONDecodeError:
                        logger.warning("收到非JSON消息: {}", data)
                        await websocket.send_text("消息格式错误，请发送JSON格式的消息")
                        continue
                    except Exception as e:
                        logger.warning("消息校验失败: {}", e)
                        await websocket.send_text(f"消息格式错误: {e}")
                        continue

                    # 处理停止指令
                    if query.action == "stop":
                        if is_generating.is_set():
                            stop_event.set()
                            logger.info("收到停止指令: session={}", query.session_id)
                            await websocket.send_text("[STOPPED]")
                        else:
                            await websocket.send_text("当前没有正在进行的生成任务")
                        continue

                    # 对话指令
                    logger.info("收到聊天查询: session_id={}", query.session_id)

                    is_generating.set()
                    stop_event.clear()
                    # 读取用户的 AI 模型配置（每次可动态切换）
                    model_config = await get_user_model_config(redis, user.id)
                    try:
                        async for chunk in chat_service.chat_query(
                            query=query,
                            stop_event=stop_event,
                            model_config=model_config,
                        ):
                            if not chunk:
                                continue
                            try:
                                await websocket.send_text(chunk)
                            except RuntimeError:
                                logger.warning("WebSocket连接已关闭，停止发送消息")
                                return
                    finally:
                        is_generating.clear()
                        stop_event.clear()

                    # 告知前端生成结束
                    try:
                        await websocket.send_text("[DONE]")
                    except RuntimeError:
                        return

                except WebSocketDisconnect:
                    logger.info("WebSocket连接已断开: {}", websocket.client)
                    return

    except CustomException as e:
        # 认证失败等业务异常
        logger.warning("WebSocket认证失败: {}", e.msg)
        await _send_error_and_close(websocket, e.msg)
    except Exception as e:
        # 未知异常
        logger.exception("WebSocket未知异常: {}", e)
        await _send_error_and_close(websocket, "服务器内部错误")