"""WebSocket 文本转发。

设备 A、B 各自带自己的 API Key 连到同一个 channel，
任一端发送的文本消息会被转发给同频道内的其他设备。

连接 URL 形如：
    ws://host/ws/{channel}?api_key=pica_xxx

API Key 放在 query 而非 header，因为浏览器 WebSocket 客户端
无法自定义请求头。
"""
from typing import Dict, Set

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

from .database import SessionLocal
from .models import ApiKey

router = APIRouter(tags=["websocket"])


class ConnectionManager:
    """按 channel 维护活跃连接，负责频道内消息转发。

    同一个 channel 下的连接互为对端：A 发的消息会推给 B（及频道内
    其他人），不会回推给 A 自己。
    """

    def __init__(self) -> None:
        # channel -> 该频道下所有活跃 WebSocket
        self._channels: Dict[str, Set[WebSocket]] = {}

    async def connect(self, channel: str, ws: WebSocket) -> None:
        await ws.accept()
        self._channels.setdefault(channel, set()).add(ws)

    def disconnect(self, channel: str, ws: WebSocket) -> None:
        conns = self._channels.get(channel)
        if not conns:
            return
        conns.discard(ws)
        if not conns:
            self._channels.pop(channel, None)

    async def broadcast(self, channel: str, sender: WebSocket, message: str) -> None:
        """把消息转发给同频道内除发送者以外的所有设备。"""
        dead: list[WebSocket] = []
        for ws in self._channels.get(channel, set()):
            if ws is sender:
                continue
            try:
                await ws.send_text(message)
            except Exception:
                # 对端已断开等异常，记下稍后清理
                dead.append(ws)
        for ws in dead:
            self.disconnect(channel, ws)

    def channel_size(self, channel: str) -> int:
        return len(self._channels.get(channel, set()))


manager = ConnectionManager()


def authenticate_api_key(db: Session, raw_key: str) -> bool:
    """校验 API Key：明文 → SHA-256 → 查库 → 命中且未吊销。"""
    if not raw_key:
        return False
    key_hash = ApiKey.hash_key(raw_key)
    api_key = db.query(ApiKey).filter(ApiKey.key_hash == key_hash).first()
    if api_key is None or api_key.revoked:
        return False
    return True


@router.websocket("/ws/{channel}")
async def websocket_endpoint(
    ws: WebSocket,
    channel: str,
    api_key: str = Query(default=""),
):
    """频道内文本转发端点。

    流程：
    1. 用 query 里的 api_key 鉴权；失败则用 1008（policy violation）关闭。
    2. 鉴权通过后 accept，加入频道。
    3. 循环收文本消息，转发给同频道其他设备。
    4. 对端断开时从频道摘除。
    """
    # 鉴权：WebSocket 阶段无法用 FastAPI 依赖注入直接返回 401，
    # 这里先校验再决定 accept / 关闭。
    db = SessionLocal()
    try:
        ok = authenticate_api_key(db, api_key)
    finally:
        db.close()

    if not ok:
        await ws.close(code=1008)  # policy violation
        return

    await manager.connect(channel, ws)
    print(
        f"[relay] join channel={channel} client={ws.client} "
        f"size={manager.channel_size(channel)}",
        flush=True,
    )
    try:
        while True:
            text = await ws.receive_text()
            print(
                f"[relay] msg channel={channel} from={ws.client} "
                f"size={manager.channel_size(channel)} msg={text!r}",
                flush=True,
            )
            await manager.broadcast(channel, ws, text)
    except WebSocketDisconnect:
        pass
    finally:
        manager.disconnect(channel, ws)
        print(
            f"[relay] leave channel={channel} client={ws.client} "
            f"size={manager.channel_size(channel)}",
            flush=True,
        )