"""Tests for the WS connect snapshot (daemon.server._handle_ws).

After the push-stream migration, a successful WS handshake sends a full
state snapshot ({type: snapshot, data: {subsystem: state|null, ...}})
instead of the retired {type: init, versions: ...} message.
"""

from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from lib.state import State


@pytest.fixture(autouse=False)
def db_reset():
    """Isolated in-memory DB so a builtin admin exists for token minting.

    Mirrors the autouse _db_reset fixture in tests/test_auth.py (the DB
    singleton must be reset and pointed at SQLite :memory: before each test).
    """
    import os

    from lib.db import get_db, reset_db_for_test

    reset_db_for_test()
    old_backend = os.environ.pop("VACUUM_WALL_DB_BACKEND", None)
    old_path = os.environ.pop("VACUUM_WALL_DB_PATH", None)
    old_seed = os.environ.pop("VACUUM_WALL_SEED_BUILTIN_ADMIN", None)

    os.environ["VACUUM_WALL_DB_BACKEND"] = "sqlite"
    os.environ["VACUUM_WALL_DB_PATH"] = ":memory:"

    get_db()  # triggers builtin-admin seed on the empty :memory: DB
    yield
    reset_db_for_test()
    if old_backend is not None:
        os.environ["VACUUM_WALL_DB_BACKEND"] = old_backend
    if old_path is not None:
        os.environ["VACUUM_WALL_DB_PATH"] = old_path
    if old_seed is not None:
        os.environ["VACUUM_WALL_SEED_BUILTIN_ADMIN"] = old_seed


@pytest.fixture
def access_token(db_reset):
    """Mint a real access token for the seeded builtin admin."""
    from lib.auth import generate_tokens

    tokens = generate_tokens("admin", {"firewall": "rw"})
    return tokens["access_token"]


class TestWsSnapshot:
    def test_snapshot_sent_on_auth_connect(self, access_token):
        """A valid JWT subprotocol yields a full snapshot after auth."""
        import asyncio

        import daemon.server as server

        store = State()
        store.set("firewall", {"zones": {"public": {}}})
        store.set("system", {"load": {"load1": 0.1}})
        # Remaining subsystems stay None (not populated).

        ws = AsyncMock()
        ws.prepare = AsyncMock()
        ws.send_json = AsyncMock()

        request = MagicMock()
        request.headers = {"Sec-WebSocket-Protocol": access_token}

        with (
            patch("aiohttp.web.WebSocketResponse", return_value=ws),
            patch.object(server, "state_store", store),
        ):
            asyncio.run(server._handle_ws(request))

        ws.send_json.assert_awaited_once()
        payload = ws.send_json.call_args[0][0]
        assert payload["type"] == "snapshot"
        data = payload["data"]
        # Every subsystem key is present (push-stream: no `init`/`versions` shape).
        for name in State.SUBSYSTEMS:
            assert name in data
        assert data["firewall"] == {"zones": {"public": {}}}
        assert data["system"] == {"load": {"load1": 0.1}}
        # Unpopulated subsystems are present but None (partial snapshot).
        assert data["dnsmasq"] is None
        assert data["wireguard"] is None

    def test_no_snapshot_without_token(self):
        """Missing token -> 401 JSON, no WS is opened, no snapshot sent."""
        import asyncio

        import daemon.server as server

        ws = AsyncMock()
        request = MagicMock()
        request.headers = {}

        with patch("aiohttp.web.WebSocketResponse") as mock_ctor:
            result = asyncio.run(server._handle_ws(request))

        assert result.status == 401
        mock_ctor.assert_not_called()
        ws.send_json.assert_not_awaited()

    def test_no_snapshot_on_invalid_token(self, access_token):
        """A token that fails validation -> 401, no snapshot sent."""
        import asyncio

        import daemon.server as server

        ws = AsyncMock()
        request = MagicMock()
        request.headers = {"Sec-WebSocket-Protocol": access_token}

        with (
            patch("aiohttp.web.WebSocketResponse") as mock_ctor,
            patch("lib.auth.validate_token", return_value=None),
        ):
            result = asyncio.run(server._handle_ws(request))

        assert result.status == 401
        mock_ctor.assert_not_called()
        ws.send_json.assert_not_awaited()