"""Abuse limits, load-based admission control and server_status messages."""

import pytest
from starlette.websockets import WebSocketDisconnect

from app.load import LoadSample

from helpers import connect, is_state, receive_until


def is_error(code):
    return lambda m: m["type"] == "error" and m.get("code") == code


def test_welcome_is_followed_by_server_status(client):
    ws = client.sockets.enter_context(client.websocket_connect("/ws"))
    ws.send_json({"type": "hello", "token": None})
    assert ws.receive_json()["type"] == "welcome"
    status = ws.receive_json()
    assert status == {"type": "server_status", "load": "normal", "accepting_games": True}


def test_message_flood_is_limited_and_closed(make_client):
    client = make_client(message_rate=0.001, message_burst=3, max_rate_violations=5)
    ws, _ = connect(client)  # uses 1 message
    for _ in range(2):  # uses the remaining burst
        ws.send_json({"type": "leave_lobby"})
    ws.send_json({"type": "leave_lobby"})  # over the limit
    assert receive_until(ws, is_error("rate_limited"))
    with pytest.raises(WebSocketDisconnect) as closed:
        for _ in range(10):
            ws.send_json({"type": "leave_lobby"})
        for _ in range(20):
            ws.receive_json()
    assert closed.value.code == 1008


def test_connections_per_ip_are_limited(make_client):
    client = make_client(max_connections_per_ip=2)
    connect(client)
    connect(client)
    ws = client.sockets.enter_context(client.websocket_connect("/ws"))
    assert ws.receive_json()["code"] == "busy"
    with pytest.raises(WebSocketDisconnect) as closed:
        ws.receive_json()
    assert closed.value.code == 1013


def test_closed_connections_free_their_slot(make_client):
    client = make_client(max_connections_per_ip=1)
    with client.websocket_connect("/ws") as ws:
        ws.send_json({"type": "hello", "token": None})
        receive_until(ws, is_state)
    ws, _ = connect(client)  # works again after the first one closed


def test_lobby_creation_rate_per_ip(make_client):
    client = make_client(lobby_create_rate=2)
    ws, _ = connect(client)
    for _ in range(2):
        ws.send_json({"type": "create_lobby", "name": "Ann"})
        receive_until(ws, lambda m: is_state(m) and m["lobby"])
    ws.send_json({"type": "create_lobby", "name": "Ann"})
    assert receive_until(ws, is_error("rate_limited"))


def test_overloaded_server_refuses_new_games(make_client):
    client = make_client(load_force_level="overloaded")
    ws = client.sockets.enter_context(client.websocket_connect("/ws"))
    ws.send_json({"type": "hello", "token": None})
    status = receive_until(ws, lambda m: m["type"] == "server_status")
    assert status["load"] == "overloaded" and not status["accepting_games"]
    ws.send_json({"type": "create_lobby", "name": "Ann"})
    assert receive_until(ws, is_error("busy"))
    assert client.get("/health").json()["accepting_games"] is False


def start_game_with_ais(client, name):
    ws, _ = connect(client)
    ws.send_json({"type": "create_lobby", "name": name})
    receive_until(ws, lambda m: is_state(m) and m["lobby"])
    ws.send_json({"type": "add_ai"})
    ws.send_json({"type": "add_ai"})
    return ws


def test_max_active_games(make_client):
    client = make_client(max_active_games=1)
    first = start_game_with_ais(client, "Ann")
    second = start_game_with_ais(client, "Ben")

    first.send_json({"type": "start_game"})
    receive_until(first, lambda m: is_state(m) and m["game"])
    # Everybody learns that no further games can start.
    status = receive_until(second, lambda m: m["type"] == "server_status")
    assert status["accepting_games"] is False

    second.send_json({"type": "start_game"})
    assert receive_until(second, is_error("busy"))

    # Leaving the running game ends it (only AIs left) and frees the slot.
    first.send_json({"type": "leave_lobby"})
    status = receive_until(second, lambda m: m["type"] == "server_status")
    assert status["accepting_games"] is True
    second.send_json({"type": "start_game"})
    assert receive_until(second, lambda m: is_state(m) and m["game"])


def test_running_game_continues_when_overloaded(client):
    ws = start_game_with_ais(client, "Ann")
    ws.send_json({"type": "start_game"})
    latest = receive_until(ws, lambda m: is_state(m) and m["game"])

    service = client.app.state.service
    for _ in range(3):  # three overloaded samples in a row
        client.portal.call(service.load.record, LoadSample(loop_lag=1.0))

    # normal -> busy (first bad sample) -> overloaded (third in a row).
    # Keep track of the latest game state seen on the way.
    while True:
        message = ws.receive_json()
        if is_state(message):
            latest = message
        if message["type"] == "server_status" and message["load"] == "overloaded":
            assert message["accepting_games"] is False
            break

    # The running game still accepts moves.
    def my_turn(state):
        obs = state["game"]["observation"]
        return obs["current_player"] == state["game"]["you"] and obs["legal_actions"]

    while not my_turn(latest):
        latest = receive_until(ws, is_state)
    ws.send_json({"type": "action", "action": latest["game"]["observation"]["legal_actions"][0]})
    reply = ws.receive_json()
    assert reply["type"] == "state", reply