import json
from io import BytesIO

import pandas as pd
import pyarrow.ipc as ipc
import pyarrow.parquet as pq
import pytest
from httpx import AsyncClient, ASGITransport
from starlette.responses import PlainTextResponse
from starlette.testclient import TestClient

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],
            }
        )

    @application.get("/health")
    async def health() -> PlainTextResponse:
        return PlainTextResponse("ok")

    return application


@pytest.fixture
def tc(app):
    with TestClient(app) as c:
        yield c


@pytest.fixture
def session_cookie(app, tc):
    """Connect via WebSocket and return the session cookie dict."""
    with tc.websocket_connect("/ws") as ws:
        msg = ws.receive_json()
        yield {SESSION_COOKIE: msg["session_id"]}


class TestHealth:
    def test_returns_ok(self, tc):
        resp = tc.get("/health")
        assert resp.status_code == 200
        assert resp.text == "ok"


class TestDemo:
    def test_default_returns_ipc_stream(self, tc, session_cookie):
        resp = tc.get("/tables/demo", cookies=session_cookie)
        assert resp.status_code == 200
        assert resp.headers["content-type"] == "application/vnd.apache.arrow.stream"
        table = ipc.open_stream(resp.content).read_all()
        assert table.num_rows == 3
        assert "id" in table.column_names
        assert "name" in table.column_names
        assert "score" in table.column_names

    def test_accept_ipc_file(self, tc, session_cookie):
        resp = tc.get(
            "/tables/demo",
            headers={"accept": "application/vnd.apache.arrow.file"},
            cookies=session_cookie,
        )
        assert resp.status_code == 200
        assert resp.headers["content-type"] == "application/vnd.apache.arrow.file"
        table = ipc.open_file(resp.content).read_all()
        assert table.num_rows == 3

    def test_accept_parquet(self, tc, session_cookie):
        resp = tc.get(
            "/tables/demo",
            headers={"accept": "application/x-parquet"},
            cookies=session_cookie,
        )
        assert resp.status_code == 200
        assert resp.headers["content-type"] == "application/x-parquet"
        table = pq.read_table(BytesIO(resp.content))
        assert table.num_rows == 3

    def test_unknown_accept_defaults_to_stream(self, tc, session_cookie):
        resp = tc.get(
            "/tables/demo",
            headers={"accept": "application/json"},
            cookies=session_cookie,
        )
        assert resp.status_code == 200
        assert resp.headers["content-type"] == "application/vnd.apache.arrow.stream"

    def test_roundtrip_data_integrity(self, tc, session_cookie):
        resp = tc.get("/tables/demo", cookies=session_cookie)
        table = ipc.open_stream(resp.content).read_all()
        assert table.column("id").to_pylist() == [1, 2, 3]
        assert table.column("name").to_pylist() == ["alice", "bob", "charlie"]
        assert table.column("score").to_pylist() == [95.5, 87.3, 91.8]


class TestTablesPushdown:
    def test_columns_projection(self, tc, session_cookie):
        resp = tc.get("/tables/demo?columns=id,name", cookies=session_cookie)
        assert resp.status_code == 200
        table = ipc.open_stream(resp.content).read_all()
        assert table.column_names == ["id", "name"]
        assert table.num_rows == 3

    def test_filter_equal(self, tc, session_cookie):
        filters = json.dumps([{"col": "id", "op": "=", "value": 1}])
        resp = tc.get(f"/tables/demo?filters={filters}", cookies=session_cookie)
        assert resp.status_code == 200
        table = ipc.open_stream(resp.content).read_all()
        assert table.num_rows == 1
        assert table.column("name").to_pylist() == ["alice"]

    def test_filter_and_columns(self, tc, session_cookie):
        filters = json.dumps([{"col": "score", "op": ">", "value": 90.0}])
        resp = tc.get(
            f"/tables/demo?columns=id,name&filters={filters}",
            cookies=session_cookie,
        )
        assert resp.status_code == 200
        table = ipc.open_stream(resp.content).read_all()
        assert table.column_names == ["id", "name"]
        assert table.num_rows == 2


class TestExceptionHandler:
    async def test_404_returns_plain_text(self):
        application = create_app()
        transport = ASGITransport(app=application)
        async with AsyncClient(transport=transport, base_url="http://test") as client:
            resp = await client.get("/nonexistent")
            assert resp.status_code == 404
            assert "application/json" not in resp.headers.get("content-type", "")
