import pandas as pd
import pyarrow.ipc as ipc
import pytest

from n6k_server.app import create_app, data
from n6k_server.session import SESSION_COOKIE


@pytest.fixture
def app():
    application = create_app()

    @data(application, "demo")
    def demo():
        return pd.DataFrame(
            {
                "id": [1, 2, 3],
                "name": ["alice", "bob", "charlie"],
                "score": [95.5, 87.3, 91.8],
            }
        )

    return application


class TestWebSocketSession:
    def test_ws_connect_returns_session_id(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws:
                msg = ws.receive_json()
                assert "session_id" in msg
                assert len(msg["session_id"]) > 0

    def test_ws_connect_registers_session(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws:
                msg = ws.receive_json()
                sid = msg["session_id"]
                assert sid in app.state.sessions

    def test_ws_disconnect_destroys_session(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws:
                msg = ws.receive_json()
                sid = msg["session_id"]
            # WebSocket closed
            assert sid not in app.state.sessions

    def test_http_with_ws_session_works(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws:
                msg = ws.receive_json()
                sid = msg["session_id"]

                resp = tc.get(
                    "/tables/demo",
                    cookies={SESSION_COOKIE: sid},
                )
                assert resp.status_code == 200
                table = ipc.open_stream(resp.content).read_all()
                assert table.num_rows == 3


class TestSessionIsolation:
    def test_two_ws_clients_get_different_sessions(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws1:
                msg1 = ws1.receive_json()
                with tc.websocket_connect("/ws") as ws2:
                    msg2 = ws2.receive_json()
                    assert msg1["session_id"] != msg2["session_id"]

    def test_view_not_visible_across_sessions(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws1:
                sid1 = ws1.receive_json()["session_id"]

                # Trigger lazy load of demo table
                tc.get("/tables/demo", cookies={SESSION_COOKIE: sid1})

                # Client 1 creates a view
                tc.post(
                    "/views/my_view",
                    content="SELECT id, name FROM demo WHERE id = 1",
                    cookies={SESSION_COOKIE: sid1},
                )

                # Client 1 can see it
                resp = tc.get("/tables/my_view", cookies={SESSION_COOKIE: sid1})
                assert resp.status_code == 200

                # Client 2 cannot
                with tc.websocket_connect("/ws") as ws2:
                    sid2 = ws2.receive_json()["session_id"]
                    resp = tc.get("/tables/my_view", cookies={SESSION_COOKIE: sid2})
                    assert resp.status_code == 404


class TestViewCreation:
    def test_create_and_fetch_view(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws:
                sid = ws.receive_json()["session_id"]
                cookies = {SESSION_COOKIE: sid}

                # Trigger lazy load of demo table
                tc.get("/tables/demo", cookies=cookies)

                resp = tc.post(
                    "/views/high_scores",
                    content="SELECT * FROM demo WHERE score > 90",
                    cookies=cookies,
                )
                assert resp.status_code == 201

                resp = tc.get("/tables/high_scores", cookies=cookies)
                assert resp.status_code == 200
                table = ipc.open_stream(resp.content).read_all()
                assert table.num_rows == 2
                names = table.column("name").to_pylist()
                assert "alice" in names
                assert "charlie" in names

    def test_create_view_bad_sql_table_400(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws:
                sid = ws.receive_json()["session_id"]
                resp = tc.post(
                    "/views/v1",
                    content="SELECT * FROM nonexistent",
                    cookies={SESSION_COOKIE: sid},
                )
                assert resp.status_code == 400

    def test_create_view_empty_body_400(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws:
                sid = ws.receive_json()["session_id"]
                resp = tc.post(
                    "/views/v1",
                    content="",
                    cookies={SESSION_COOKIE: sid},
                )
                assert resp.status_code == 400

    def test_create_view_non_select_rejected(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws:
                sid = ws.receive_json()["session_id"]
                resp = tc.post(
                    "/views/v1",
                    content="DROP TABLE demo",
                    cookies={SESSION_COOKIE: sid},
                )
                assert resp.status_code == 400

    def test_create_view_invalid_sql_400(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws:
                sid = ws.receive_json()["session_id"]
                resp = tc.post(
                    "/views/v1",
                    content="SELECT * FROM nonexistent_table",
                    cookies={SESSION_COOKIE: sid},
                )
                assert resp.status_code == 400

    def test_view_id_validation(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws:
                sid = ws.receive_json()["session_id"]
                resp = tc.post(
                    "/views/invalid-id",
                    content="SELECT 1",
                    cookies={SESSION_COOKIE: sid},
                )
                assert resp.status_code == 400

    def test_replace_view(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws:
                sid = ws.receive_json()["session_id"]
                cookies = {SESSION_COOKIE: sid}

                # Trigger lazy load of demo table
                tc.get("/tables/demo", cookies=cookies)

                tc.post(
                    "/views/v1",
                    content="SELECT * FROM demo WHERE id = 1",
                    cookies=cookies,
                )

                resp = tc.post(
                    "/views/v1",
                    content="SELECT * FROM demo WHERE id = 2",
                    cookies=cookies,
                )
                assert resp.status_code == 201

                resp = tc.get("/tables/v1", cookies=cookies)
                table = ipc.open_stream(resp.content).read_all()
                assert table.num_rows == 1
                assert table.column("id").to_pylist() == [2]


class TestBaseTableReadOnly:
    def test_insert_into_base_table_fails(self, app):
        from starlette.testclient import TestClient

        with TestClient(app) as tc:
            with tc.websocket_connect("/ws") as ws:
                sid = ws.receive_json()["session_id"]
                resp = tc.post(
                    "/views/bad",
                    content="SELECT * FROM (INSERT INTO demo VALUES (4, 'dave', 88.0))",
                    cookies={SESSION_COOKIE: sid},
                )
                assert resp.status_code == 400
