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 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.
            assert admin_ws.receive_json()["type"] == "state"
            assert admin_ws.receive_json()["type"] == "lobby_update"

            # First "next" opens Q0 from lobby.
            admin_ws.send_json({"type": "next"})
            assert student_ws.receive_json()["type"] == "question_open"
            admin_open = admin_ws.receive_json()
            assert admin_open["type"] == "question_open"
            assert admin_ws.receive_json()["type"] == "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:
            assert admin_ws.receive_json()["type"] == "state"
            assert admin_ws.receive_json()["type"] == "lobby_update"
            admin_ws.send_json({"type": "open_question", "question_idx": 0, "time_limit": 2})
            assert student_ws.receive_json()["type"] == "question_open"
            assert admin_ws.receive_json()["type"] == "question_open"
            assert admin_ws.receive_json()["type"] == "live_histogram"

            admin_ws.send_json({"type": "close_question"})
            assert student_ws.receive_json()["type"] == "question_closed"
            admin_msgs = [admin_ws.receive_json(), admin_ws.receive_json()]
            assert {m["type"] for m in admin_msgs} == {"question_closed", "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:
        assert admin_ws.receive_json()["type"] == "state"
        assert admin_ws.receive_json()["type"] == "lobby_update"
        admin_ws.send_json({"type": "open_question", "question_idx": 0, "time_limit": 2})
        assert admin_ws.receive_json()["type"] == "question_open"
        assert admin_ws.receive_json()["type"] == "live_histogram"

        admin_ws.send_json({"type": "reset"})
        # After reset, the instructor receives a state=lobby snapshot + lobby_update.
        msgs = []
        while len(msgs) < 2:
            msgs.append(admin_ws.receive_json())
        types = [m["type"] for m in msgs]
        assert "state" in types