import json
from io import BytesIO

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

import n6k_server


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

    @n6k_server.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 transport(app):
    return ASGITransport(app=app)


@pytest.fixture
async def client(transport):
    async with AsyncClient(transport=transport, base_url="http://test") as c:
        yield c


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


class TestDemo:
    async def test_default_returns_ipc_stream(self, client: AsyncClient):
        resp = await client.get("/tables/demo")
        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

    async def test_accept_ipc_file(self, client: AsyncClient):
        resp = await client.get(
            "/tables/demo",
            headers={"accept": "application/vnd.apache.arrow.file"},
        )
        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

    async def test_accept_parquet(self, client: AsyncClient):
        resp = await client.get(
            "/tables/demo",
            headers={"accept": "application/x-parquet"},
        )
        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

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

    async def test_roundtrip_data_integrity(self, client: AsyncClient):
        resp = await client.get("/tables/demo")
        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:
    async def test_columns_projection(self, client: AsyncClient):
        resp = await client.get("/tables/demo?columns=id,name")
        assert resp.status_code == 200
        table = ipc.open_stream(resp.content).read_all()
        assert table.column_names == ["id", "name"]
        assert table.num_rows == 3

    async def test_filter_equal(self, client: AsyncClient):
        filters = json.dumps([{"col": "id", "op": "=", "value": 1}])
        resp = await client.get(f"/tables/demo?filters={filters}")
        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"]

    async def test_filter_and_columns(self, client: AsyncClient):
        filters = json.dumps([{"col": "score", "op": ">", "value": 90.0}])
        resp = await client.get(f"/tables/demo?columns=id,name&filters={filters}")
        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, client: AsyncClient):
        resp = await client.get("/nonexistent")
        assert resp.status_code == 404
        assert "application/json" not in resp.headers.get("content-type", "")
