"""WebSocket 路由：流式对话 + 重新生成。协议：
客户端 → {"type":"chat","content":"...","agent_id":null,"model":"","file_ids":[]}
        {"type":"regenerate"}
服务端 → {"type":"delta","message_id":1,"kind":"reasoning"|"content","content":"增量文本"}
        {"type":"done","message_id":1}
        {"type":"suggestions","message_id":1,"items":["..",".."]}
        {"type":"title","title":"..."}
        {"type":"error","message":"..."}
"""
import asyncio

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from sqlalchemy.orm import Session

from ..core.security import decode_token
from ..database import SessionLocal
from ..models import ChatSession, User
from ..services import chat_service

router = APIRouter()

# session_id -> set[WebSocket]（用于推送建议/标题）
_conns: dict[int, set[WebSocket]] = {}


def push_to_session(session_id: int, payload: dict):
    """向该会话的所有在线连接推送消息。"""
    for ws in list(_conns.get(session_id, set())):
        try:
            asyncio.get_event_loop().create_task(ws.send_json(payload))
        except Exception:
            pass


@router.websocket("/chat/ws")
async def chat_ws(websocket: WebSocket):
    await websocket.accept()
    token = websocket.query_params.get("token", "")
    session_id = websocket.query_params.get("session_id", "")
    payload = decode_token(token)
    if not payload:
        await websocket.send_json({"type": "error", "message": "认证失败"})
        await websocket.close()
        return

    db: Session = SessionLocal()
    try:
        user = db.get(User, int(payload["sub"]))
        if not user or not user.is_active:
            await websocket.send_json({"type": "error", "message": "用户无效"})
            await websocket.close()
            return

        session = db.get(ChatSession, int(session_id)) if session_id.isdigit() else None
        if not session or session.user_id != user.id:
            await websocket.send_json({"type": "error", "message": "会话不存在"})
            await websocket.close()
            return

        _conns.setdefault(session.id, set()).add(websocket)
        try:
            while True:
                data = await websocket.receive_json()
                msg_type = data.get("type")

                if msg_type == "chat":
                    content = (data.get("content") or "").strip()
                    file_ids = data.get("file_ids") or []
                    if not content and not file_ids:
                        await websocket.send_json({"type": "error", "message": "消息不能为空"})
                        continue
                    try:
                        async for (kind, delta), message_id in chat_service.chat_stream(
                            db, user, session, content,
                            agent_id=data.get("agent_id"), model=data.get("model", ""),
                            file_ids=file_ids,
                        ):
                            await websocket.send_json({"type": "delta", "message_id": message_id, "kind": kind, "content": delta})
                        await websocket.send_json({"type": "done", "message_id": message_id})
                    except Exception as e:
                        await websocket.send_json({"type": "error", "message": f"模型调用失败：{e}"})

                elif msg_type == "regenerate":
                    try:
                        async for (kind, delta), message_id in chat_service.regenerate_stream(db, session):
                            await websocket.send_json({"type": "delta", "message_id": message_id, "kind": kind, "content": delta})
                        await websocket.send_json({"type": "done", "message_id": message_id})
                    except LookupError as e:
                        await websocket.send_json({"type": "error", "message": str(e)})
                    except Exception as e:
                        await websocket.send_json({"type": "error", "message": f"模型调用失败：{e}"})
        finally:
            _conns.get(session.id, set()).discard(websocket)
            if not _conns.get(session.id):
                _conns.pop(session.id, None)
    except WebSocketDisconnect:
        pass
    finally:
        db.close()