import asyncio

import pytest
from websockets import connect as ws_connect
from websockets import serve
from websockets.exceptions import ConnectionClosed

from quantumpipe.exceptions import ProtocolError
from quantumpipe.sansio import (
    _CREDENTIAL,
    Message,
    Verification,
    generate_registration_secret,
)
from quantumpipe.ws import TEXT, SecSocket

AD = b"test app"


class FakeWS:
    """websockets-library style endpoint (send/recv) over memory queues."""

    def __init__(self, inbox, outbox):
        self._inbox = inbox
        self._outbox = outbox

    async def send(self, data):
        self._outbox.put_nowait(data)

    async def recv(self):
        data = await self._inbox.get()
        if data is None:  # the peer closed
            raise ConnectionClosed(None, None)
        return data

    async def close(self):
        self._outbox.put_nowait(None)  # wake the peer like a real close


class FakeStarletteWS(FakeWS):
    """FastAPI/Starlette style endpoint (accept/send_bytes/receive_bytes)."""

    accepted = False

    async def accept(self):
        self.accepted = True

    async def send_bytes(self, data):
        self._outbox.put_nowait(data)

    async def receive_bytes(self):
        return await self.recv()


def ws_pair(starlette=False):
    """A connected pair of in-memory WebSocket endpoints (srv, cli)."""
    to_srv, to_cli = asyncio.Queue(), asyncio.Queue()
    cls = FakeStarletteWS if starlette else FakeWS
    return cls(to_srv, to_cli), cls(to_cli, to_srv)


async def socket_pair(
    cred_srv, cred_cli, verify_srv=None, verify_cli=None, ad=AD, starlette=False
):
    """Two SecSockets handshake over a fake ws pair; returns (srv, cli)."""
    ws_srv, ws_cli = ws_pair(starlette)
    srv = SecSocket(
        ws_srv,
        {"x": cred_srv},
        ad=ad,
        handlers={Verification: verify_srv} if verify_srv else None,
    )
    cli = SecSocket(
        ws_cli,
        {"x": cred_cli},
        "x",
        ad=ad,
        handlers={Verification: verify_cli} if verify_cli else None,
    )
    await asyncio.gather(srv.accept(), cli.connect())
    return srv, cli


async def test_roundtrip():
    psk = generate_registration_secret()
    srv, cli = await socket_pair(psk, psk)
    await cli.send(b"ping")
    assert await srv.receive() == b"ping"
    await srv.send(b"pong")
    assert await cli.receive() == b"pong"


async def test_send_picks_frame_type():
    # str goes as a TEXT frame and is delivered decoded; bytes are opaque
    psk = generate_registration_secret()
    srv, cli = await socket_pair(psk, psk)
    await cli.send("hello")
    assert await srv.receive() == "hello"
    await srv.send(b"\x00\xff")
    assert await cli.receive() == b"\x00\xff"
    await cli.send("")
    assert await srv.receive() == ""  # empty text frame
    await cli.send(b"")
    assert await srv.receive() == b""  # empty binary frame


async def test_message_passthrough():
    # applications can use their own kinds (protobuf etc.): a Message
    # passes through the wrapper untouched
    psk = generate_registration_secret()
    srv, cli = await socket_pair(psk, psk)
    await cli.send(Message(b"\x89fake-protobuf", 7))
    assert await srv.receive() == Message(b"\x89fake-protobuf", 7)


async def test_starlette_flavor():
    # the FastAPI/Starlette duck-type: accept + send_bytes/receive_bytes
    psk = generate_registration_secret()
    ws_srv, ws_cli = ws_pair(starlette=True)
    srv = SecSocket(ws_srv, {"x": psk}, ad=AD)
    cli = SecSocket(ws_cli, {"x": psk}, "x", ad=AD)
    await asyncio.gather(srv.accept(), cli.connect())
    assert ws_srv.accepted
    await cli.send(b"ping")
    assert await srv.receive() == b"ping"


async def test_protocol_kind_rejected():
    # kind 0 is reserved for in-band protocol messages; none implemented
    psk = generate_registration_secret()
    srv, cli = await socket_pair(psk, psk)
    with pytest.raises(ProtocolError, match="reserved"):
        await cli.send(Message(b"x", 0))
    frame = bytes(cli.session._tx.encrypt(bytes([0]) + b"x"))
    await cli._ws.send(frame)  # bypass the wrapper
    with pytest.raises(ProtocolError, match="reserved"):
        await srv.receive()


async def test_empty_frame_rejected():
    # a box with no plaintext at all carries no kind byte
    psk = generate_registration_secret()
    srv, cli = await socket_pair(psk, psk)
    await cli._ws.send(bytes(cli.session._tx.encrypt(b"")))
    with pytest.raises(ProtocolError, match="empty"):
        await srv.receive()


async def test_bad_utf8_text_frame_rejected():
    psk = generate_registration_secret()
    srv, cli = await socket_pair(psk, psk)
    (frame,) = cli.session(Message(b"\xff\xff", TEXT))
    await cli._ws.send(frame)
    with pytest.raises(ProtocolError, match="utf-8"):
        await srv.receive()


async def test_traffic_before_establishment_rejected():
    psk = generate_registration_secret()
    _ws_srv, ws_cli = ws_pair()
    cli = SecSocket(ws_cli, {"x": psk}, "x", ad=AD)
    with pytest.raises(ProtocolError, match="no traffic"):
        await cli.send(b"early")


async def test_verification_callback():
    psk = generate_registration_secret()
    seen = []

    async def skip(v: Verification):  # -> None: skip the comparison
        seen.append(v.sas)

    srv, cli = await socket_pair(psk, psk, skip, skip)
    cred_srv, cred_cli = srv.session.credential, cli.session.credential
    assert len(seen) == 2 and seen[0] == seen[1]
    # skipped comparison: established, but flagged unverified
    assert not srv.session.verified and not cli.session.verified
    assert cli.session.verification.sas == seen[1]

    # the resumption flow never prompts: it is silent even while
    # unverified; the stored digest is exposed and the parties confirm
    # whenever they wish
    seen.clear()
    srv, cli = await socket_pair(cred_srv, cred_cli, skip, skip)
    await cli.send(b"resumed")
    assert await srv.receive() == b"resumed"
    assert seen == []  # no comparison asked
    assert not srv.session.verified and not cli.session.verified
    assert cli.session.verification.sas == _CREDENTIAL.unpack(cred_cli)[4]
    cli.session.confirm(True)  # compared out of band: mark verified
    srv.session.confirm(True)
    assert cli.session.verified and srv.session.verified
    cred_srv, cred_cli = srv.session.credential, cli.session.credential

    # a verified credential resumes silently, flagged verified
    srv, cli = await socket_pair(cred_srv, cred_cli, skip, skip)
    await cli.send(b"again")
    assert await srv.receive() == b"again"
    assert seen == []
    assert srv.session.verified and cli.session.verified


async def test_verification_rejected():
    psk = generate_registration_secret()

    async def reject(v: Verification):
        return False

    ws_srv, ws_cli = ws_pair()
    srv_store = {"x": psk}
    srv = SecSocket(ws_srv, srv_store, ad=AD)
    cli_store = {"x": psk}
    cli = SecSocket(ws_cli, cli_store, "x", ad=AD, handlers={Verification: reject})
    srv_task = asyncio.create_task(srv.accept())
    # the responder's default answer releases msg4, but the initiator's
    # reject holds back msg5: her side fails, the responder never swaps
    with pytest.raises(ProtocolError, match="rejected"):
        await cli.connect()
    assert cli.session.state == "failed"
    assert cli_store == {}  # rejected: the credential is deleted
    with pytest.raises(ConnectionClosed):
        await srv_task
    assert srv.session.state == "reg:msg5"  # parked: no msg5 ever came
    assert srv_store == {"x": psk}  # the registration key is intact


async def test_parallel_registration_race():
    """Two registrations racing on one registration key: the swap at the
    comparison answer arbitrates -- the first answer (the app default
    skip, right after msg3) wins the store slot, the loser's server
    session fails with "registration key already spent". Both clients
    still establish (the loser's composite is orphaned server-side)."""
    psk = generate_registration_secret()
    srv_store = {"x": psk}
    outcomes = []

    async def server_side(ws):
        sock = SecSocket(ws, srv_store, ad=AD)
        try:
            await sock.accept()
        except Exception as e:  # noqa: BLE001 -- recorded and asserted below
            outcomes.append(e)
        else:
            outcomes.append("established")

    async def client_side():
        ws_srv, ws_cli = ws_pair()
        srv_task = asyncio.create_task(server_side(ws_srv))
        store = {"x": psk}
        sock = SecSocket(ws_cli, store, "x", ad=AD)
        await sock.connect()
        assert sock.session.state == "established"
        return srv_task, store

    (t1, s1), (t2, s2) = await asyncio.gather(client_side(), client_side())
    await asyncio.gather(t1, t2)

    # exactly one winner server-side; the loser is rejected at the swap
    fails = [o for o in outcomes if isinstance(o, ProtocolError)]
    assert outcomes.count("established") == 1
    assert len(fails) == 1 and "registration key already spent" in str(fails[0])
    # the winner's composite replaced the registration key; the loser's
    # burn/delete machinery never touched the slot again
    assert len(srv_store) == 1 and len(srv_store["x"]) == _CREDENTIAL.size
    # both clients rotated their own stores and established
    assert all(len(s["x"]) == _CREDENTIAL.size for s in (s1, s2))


async def test_real_websockets_roundtrip():
    """Integration over a real localhost websockets server/client."""
    psk = generate_registration_secret()
    srv_sock = None

    async def handler(ws):
        nonlocal srv_sock
        srv_sock = SecSocket(ws, {"x": psk}, ad=AD)
        await srv_sock.accept()
        await ws.wait_closed()

    server = await serve(handler, "127.0.0.1", 0)
    port = server.sockets[0].getsockname()[1]
    ws = await ws_connect(f"ws://127.0.0.1:{port}")
    cli = SecSocket(ws, {"x": psk}, "x", ad=AD)
    await cli.connect()
    while srv_sock is None or srv_sock.session.state != "established":
        await asyncio.sleep(0.01)
    await cli.send("ping")
    assert await srv_sock.receive() == "ping"
    await srv_sock.send(b"pong")
    assert await cli.receive() == b"pong"
    await ws.close()
    server.close()