import pyarrow as pa
import pyarrow.ipc as ipc
import pyarrow.parquet as pq
from io import BytesIO
from starlette.testclient import TestClient

from n6k_server.formats import ArrowFormat, negotiate, arrow_response, _serialize_ipc_stream, _serialize_ipc_file, _serialize_parquet


def _make_table() -> pa.Table:
    return pa.table({"x": [1, 2, 3], "y": ["a", "b", "c"]})


class _FakeRequest:
    def __init__(self, accept: str = ""):
        self.headers = {"accept": accept}


class TestNegotiate:
    def test_ipc_stream(self):
        req = _FakeRequest("application/vnd.apache.arrow.stream")
        assert negotiate(req) == ArrowFormat.IPC_STREAM

    def test_ipc_file(self):
        req = _FakeRequest("application/vnd.apache.arrow.file")
        assert negotiate(req) == ArrowFormat.IPC_FILE

    def test_parquet(self):
        req = _FakeRequest("application/x-parquet")
        assert negotiate(req) == ArrowFormat.PARQUET

    def test_default_when_missing(self):
        req = _FakeRequest("")
        assert negotiate(req) == ArrowFormat.IPC_STREAM

    def test_default_when_unknown(self):
        req = _FakeRequest("application/json")
        assert negotiate(req) == ArrowFormat.IPC_STREAM

    def test_first_match_wins(self):
        req = _FakeRequest("application/x-parquet, application/vnd.apache.arrow.file")
        assert negotiate(req) == ArrowFormat.PARQUET

    def test_quality_params_ignored_gracefully(self):
        req = _FakeRequest("application/vnd.apache.arrow.file;q=0.9")
        assert negotiate(req) == ArrowFormat.IPC_FILE


class TestSerializers:
    def test_ipc_stream_roundtrip(self):
        table = _make_table()
        data = _serialize_ipc_stream(table)
        result = ipc.open_stream(data).read_all()
        assert result.equals(table)

    def test_ipc_file_roundtrip(self):
        table = _make_table()
        data = _serialize_ipc_file(table)
        result = ipc.open_file(data).read_all()
        assert result.equals(table)

    def test_parquet_roundtrip(self):
        table = _make_table()
        data = _serialize_parquet(table)
        result = pq.read_table(BytesIO(data))
        assert result.equals(table)


class TestArrowResponse:
    def test_content_type_ipc_stream(self):
        table = _make_table()
        resp = arrow_response(table, ArrowFormat.IPC_STREAM)
        assert resp.media_type == "application/vnd.apache.arrow.stream"

    def test_content_type_parquet(self):
        table = _make_table()
        resp = arrow_response(table, ArrowFormat.PARQUET)
        assert resp.media_type == "application/x-parquet"

    def test_body_is_deserializable(self):
        table = _make_table()
        resp = arrow_response(table, ArrowFormat.IPC_STREAM)
        result = ipc.open_stream(resp.body).read_all()
        assert result.equals(table)
