import pytest
from starlette.websockets import WebSocketDisconnect

from conftest import admin_login, create_session, join_student


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


def test_instructor_controls_transition_and_broadcast(client, sample_pool):
    sid = create_session(client, sample_pool)
    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})
            student_open = student_ws.receive_json()
            assert student_open["type"] == "question_open"
            admin_open = admin_ws.receive_json()
            assert admin_open["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"
            messages = [admin_ws.receive_json(), admin_ws.receive_json()]
            assert {msg["type"] for msg in messages} == {"question_closed", "full_leaderboard"}

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