import json
import logging

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from app.engine.models import ConversationSource, ConversationStatus
from app.engine.routes.runtime import get_runtime
from app.engine.services.conn_limiter import DEFAULT_WS_CONN_LIMIT
from app.engine.services.conn_limiter import limiter
from app.engine.services.conversation import ConversationManager

logger = logging.getLogger(__name__)
router = APIRouter()


@router.websocket("/ws/chat")
async def ws_chat(websocket: WebSocket):
    # Task 8: 每 (租户, IP) 并发连接上限——必须在 accept 前拒绝(close 代
    # accept,客户端握手即断)。IP 提取与下方 client_info 同一来源(XFF 优先,
    # scope client 兜底),握手期 headers/scope 已就绪、不依赖首条 init 消息;
    # 客户端自报的 client_info.ip 仍按原逻辑只进会话记录,不参与限流。
    runtime = get_runtime(websocket)
    client_ip = websocket.client.host if websocket.client else ""
    forwarded = websocket.headers.get("x-forwarded-for", "").split(",")[0].strip()
    if forwarded:
        client_ip = forwarded

    limit = runtime.config.data["engine"].get(
        "ws_conn_limit", DEFAULT_WS_CONN_LIMIT)
    if not limiter.try_acquire(runtime.config.slug, client_ip, limit=limit):
        await websocket.close(code=1013, reason="too many connections")
        return

    session_id = None

    try:
        await websocket.accept()
        # Task 5b: 租户运行时经 state/默认租户兜底解析(过渡期无中间件);
        # manager 由装配线从 runtime 造(db 会话 + runtime 引用)
        async with runtime.session_factory() as db:
            manager = ConversationManager(db, runtime=runtime)

            # Collect client info(client_ip/forwarded 已在握手期提取,复用)
            user_agent = websocket.headers.get("user-agent", "")
            client_info = {"ip": client_ip, "user_agent": user_agent}

            # First message: init or resume
            init_data = await websocket.receive_text()
            init_msg = json.loads(init_data)

            # Task 10: 客户端自报 site(空域名站点如 es 靠嵌入代码 data-site
            # 接入)作会话站点标——只加记录,不参与业务逻辑;写入点在下方
            # 新建会话处(优先级对齐 REST,见那里注释)
            site_slug = str(init_msg.get("site") or "").strip()

            if "session_id" in init_msg:
                conv = await manager.get_by_session(init_msg["session_id"])
                if conv and conv.status == ConversationStatus.active:
                    session_id = conv.session_id
                    await websocket.send_text(json.dumps({
                        "type": "resumed",
                        "session_id": session_id,
                    }))
                else:
                    session_id = None

            # page_token: reuse active conversation from same page session
            if not session_id and "page_token" in init_msg and init_msg["page_token"]:
                conv = await manager.get_by_page_token(init_msg["page_token"])
                if conv:
                    session_id = conv.session_id
                    await websocket.send_text(json.dumps({
                        "type": "resumed",
                        "session_id": session_id,
                    }))

            if not session_id:
                source = ConversationSource(
                    init_msg.get("source", "standalone")
                )
                merged_info = {**client_info, **(init_msg.get("client_info") or {})}
                # 服务端 XFF 优先于客户端自报 ip（客户端自报可伪造，可伪造面最小）；
                # 无 XFF 时保留客户端自报，user_agent 等其他字段不受影响。
                # （REST 侧已同步上游用直连 host，WS 保持自报——见
                #   tests/test_chat_rest_ip.py）
                if forwarded:
                    merged_info["ip"] = forwarded
                # Task 10 复审: 专用 site 字段同样赢过 init_msg.client_info
                # 自报的 site(与 REST ?site= 覆盖 req.client_info 的优先级
                # 对齐——专用通道是嵌入代码显式声明,可信度高于任意自报 dict)
                if site_slug:
                    merged_info["site"] = site_slug
                conv = await manager.create_conversation(
                    source=source,
                    client_info=merged_info,
                    page_token=init_msg.get("page_token"),
                )
                session_id = conv.session_id
                await websocket.send_text(json.dumps({
                    "type": "init",
                    "session_id": session_id,
                }))

            # Chat loop
            while True:
                data = await websocket.receive_text()
                msg = json.loads(data)
                content = msg.get("content", "").strip()
                if not content:
                    continue

                ended = False

                async for event in manager.stream_process_message(session_id, content):
                    if event["type"] == "error":
                        await websocket.send_text(json.dumps({
                            "type": "error",
                            "message": event["message"],
                        }))
                        break

                    await websocket.send_text(json.dumps(event, ensure_ascii=False))

                    if event["type"] == "evaluated" and event.get("ended"):
                        ended = True
                        break

                if ended:
                    break

    except WebSocketDisconnect:
        logger.info("Client disconnected: %s", session_id)
    except Exception:
        logger.exception("WebSocket error for session %s", session_id)
    finally:
        # Task 8: 正常关闭与异常断开都归还名额——acquire 成功后所有路径
        # 必经此处,不泄漏(accept 失败/循环中抛错同样覆盖)
        limiter.release(runtime.config.slug, client_ip)