from enum import StrEnum
from io import BytesIO

import pyarrow as pa
import pyarrow.ipc as ipc
import pyarrow.parquet as pq
from starlette.requests import Request
from starlette.responses import Response


class ArrowFormat(StrEnum):
    IPC_STREAM = "application/vnd.apache.arrow.stream"
    IPC_FILE = "application/vnd.apache.arrow.file"
    PARQUET = "application/x-parquet"


_ACCEPT_MAP: dict[str, ArrowFormat] = {f.value: f for f in ArrowFormat}


def negotiate(request: Request) -> ArrowFormat:
    accept = request.headers.get("accept", "")
    for media_type in accept.split(","):
        key = media_type.strip().split(";")[0].strip()
        if key in _ACCEPT_MAP:
            return _ACCEPT_MAP[key]
    return ArrowFormat.IPC_STREAM


def _serialize_ipc_stream(table: pa.Table) -> bytes:
    sink = pa.BufferOutputStream()
    with ipc.new_stream(sink, table.schema) as writer:
        writer.write_table(table)
    return sink.getvalue().to_pybytes()


def _serialize_ipc_file(table: pa.Table) -> bytes:
    sink = pa.BufferOutputStream()
    with ipc.new_file(sink, table.schema) as writer:
        writer.write_table(table)
    return sink.getvalue().to_pybytes()


def _serialize_parquet(table: pa.Table) -> bytes:
    buf = BytesIO()
    pq.write_table(table, buf)
    return buf.getvalue()


_SERIALIZERS = {
    ArrowFormat.IPC_STREAM: _serialize_ipc_stream,
    ArrowFormat.IPC_FILE: _serialize_ipc_file,
    ArrowFormat.PARQUET: _serialize_parquet,
}


def arrow_response(table: pa.Table, fmt: ArrowFormat) -> Response:
    body = _SERIALIZERS[fmt](table)
    return Response(content=body, media_type=fmt.value)
