# Task 8: WS 每 IP 并发连接上限(spec §6.1 P2)。
#
# 两层测试:
# - 纯单元: ConnectionLimiter 算法(名额/键隔离/归零清 key/no-op release);
# - 接入: /ws/chat 握手期按 (租户, IP) 限流——第 limit+1 条并发连接被
#   1013 拒(close 代 accept),连接关闭后名额必须释放。
import json
from contextlib import ExitStack

import pytest
from starlette.websockets import WebSocketDisconnect

from app.engine.services.conn_limiter import ConnectionLimiter

# ---- 纯单元: 算法 -----------------------------------------------------------


def test_limit_per_ip():
    lim = ConnectionLimiter(limit=2)
    assert lim.try_acquire("t1", "1.1.1.1") and lim.try_acquire("t1", "1.1.1.1")
    assert not lim.try_acquire("t1", "1.1.1.1")          # 第 3 个拒绝
    lim.release("t1", "1.1.1.1")
    assert lim.try_acquire("t1", "1.1.1.1")              # 释放后可再进
    assert lim.try_acquire("t2", "1.1.1.1")              # 其他租户不受影响
    assert lim.try_acquire("t1", "2.2.2.2")              # 其他 IP 不受影响


def test_count_zero_cleans_key():
    # 计数归零删 key: 长期运行下内部 dict 不随历史 (租户,IP) 对无限增长
    lim = ConnectionLimiter(limit=3)
    lim.try_acquire("t", "1.1.1.1")
    lim.try_acquire("t", "1.1.1.1")
    lim.release("t", "1.1.1.1")
    assert ("t", "1.1.1.1") in lim._counts               # 还剩 1 个,键保留
    lim.release("t", "1.1.1.1")
    assert lim._counts == {}                              # 归零即清空


def test_release_unheld_key_is_safe_noop():
    lim = ConnectionLimiter(limit=2)
    lim.release("nobody", "9.9.9.9")                      # 从未 acquire 过
    lim.try_acquire("t", "1.1.1.1")
    lim.release("t", "1.1.1.1")
    lim.release("t", "1.1.1.1")                           # 双重 release
    assert lim._counts == {}                              # 不出负计数
    assert lim.try_acquire("t", "1.1.1.1")                # 状态仍可用


# ---- 接入: /ws/chat 握手期限流 ----------------------------------------------
# TestClient 的 WS scope client 恒为 ("testclient", 50000) 且测试不带 XFF,
# 故所有连接同落 ip="testclient";app_client 默认 jiancai 域——同租户同 IP,
# 正好是被限的键。默认 ws_conn_limit=10(defaults.py engine 组)。


def test_ws_over_limit_rejected_with_1013(app_client):
    """填满 10 个名额 → 第 11 条握手即被拒(1013, 未 accept)。"""
    with ExitStack() as stack:
        conns = [
            stack.enter_context(app_client.websocket_connect("/ws/chat"))
            for _ in range(10)
        ]
        # 前 10 条真实可用(走完 init),不是恰好没崩
        conns[0].send_text(json.dumps({"source": "standalone"}))
        assert conns[0].receive_json()["type"] == "init"

        with pytest.raises(WebSocketDisconnect) as ei:
            with app_client.websocket_connect("/ws/chat"):
                pass  # pragma: no cover — 未 accept 即 close,进不来
        assert ei.value.code == 1013
        assert ei.value.reason == "too many connections"


def test_ws_conn_release_allows_reentry(app_client):
    with ExitStack() as stack:
        for _ in range(10):
            stack.enter_context(app_client.websocket_connect("/ws/chat"))
    # ExitStack 退出 = 10 条全部关闭,服务端 finally 已 release
    with app_client.websocket_connect("/ws/chat") as ws:
        ws.send_text(json.dumps({"source": "standalone"}))
        assert ws.receive_json()["type"] == "init"