import pytest
from starlette.websockets import WebSocketDisconnect

from conftest import admin_login, join_student


def test_instructor_ws_requires_admin_cookie(client, sid):
    with pytest.raises(WebSocketDisconnect) as exc:
        with client.websocket_connect(f"/ws/instructor/{sid}"):
            pass
    assert exc.value.code == 4001


def _drain_until(ws, target_type, max_msgs=12):
    """Helper: pull messages off `ws` until one matches `target_type`. Lets
    tests skip auxiliary state-tracking messages (presence_update,
    full_leaderboard) that fire as side-effects of state changes."""
    for _ in range(max_msgs):
        msg = ws.receive_json()
        if msg["type"] == target_type:
            return msg
    raise AssertionError(f"did not see message {target_type!r} after {max_msgs} attempts")


def test_instructor_next_command_drives_full_loop(client, sid):
    """The 'next' WS message drives the entire lifecycle:
    lobby → opens Q0 → closes Q0 + opens Q1 → ... → closes last + ends."""
    join_student(client, sid, "s1", "Student One")
    admin_login(client)
    with client.websocket_connect(f"/ws/student/{sid}") as student_ws:
        assert student_ws.receive_json()["type"] == "state"
        with client.websocket_connect(f"/ws/instructor/{sid}") as admin_ws:
            # Drain lobby snapshot (state + lobby_update + presence_update).
            _drain_until(admin_ws, "presence_update")

            # First "next" opens Q0 from lobby.
            admin_ws.send_json({"type": "next"})
            assert student_ws.receive_json()["type"] == "question_open"
            _drain_until(admin_ws, "question_open")
            _drain_until(admin_ws, "live_histogram")

            # Second "next" closes Q0 and opens Q1.
            admin_ws.send_json({"type": "next"})
            student_msgs = [student_ws.receive_json() for _ in range(2)]
            assert {m["type"] for m in student_msgs} == {"question_closed", "question_open"}


def test_instructor_close_then_next_emits_clean_open(client, sid):
    join_student(client, sid, "s1", "Student One")
    admin_login(client)
    with client.websocket_connect(f"/ws/student/{sid}") as student_ws:
        assert student_ws.receive_json()["type"] == "state"
        with client.websocket_connect(f"/ws/instructor/{sid}") as admin_ws:
            _drain_until(admin_ws, "presence_update")
            admin_ws.send_json({"type": "open_question", "question_idx": 0, "time_limit": 2})
            assert student_ws.receive_json()["type"] == "question_open"
            _drain_until(admin_ws, "question_open")
            _drain_until(admin_ws, "live_histogram")

            admin_ws.send_json({"type": "close_question"})
            assert student_ws.receive_json()["type"] == "question_closed"
            _drain_until(admin_ws, "question_closed")
            _drain_until(admin_ws, "full_leaderboard")

            admin_ws.send_json({"type": "next"})
            assert student_ws.receive_json()["type"] == "question_open"


def test_reset_command_returns_session_to_lobby(client, sid):
    join_student(client, sid, "s1", "Student One")
    admin_login(client)
    with client.websocket_connect(f"/ws/instructor/{sid}") as admin_ws:
        _drain_until(admin_ws, "presence_update")
        admin_ws.send_json({"type": "open_question", "question_idx": 0, "time_limit": 2})
        _drain_until(admin_ws, "question_open")
        _drain_until(admin_ws, "live_histogram")

        admin_ws.send_json({"type": "reset"})
        # After reset, instructor receives a state=lobby snapshot.
        msg = _drain_until(admin_ws, "state")
        assert msg["state"] == "lobby"