import pandas as pd
import pyarrow as pa
from fastapi import APIRouter, Request
from starlette.responses import Response

from n6k_server.formats import arrow_response, negotiate
from n6k_server.pushdown import apply_columns, apply_filters

router = APIRouter()


@router.get("/")
async def list_tables(request: Request) -> Response:
    names = list(request.app.state.data_sources.keys())
    table = pa.table({
        "name": pa.array(names, type=pa.utf8()),
        "path": pa.array([f"/tables/{n}" for n in names], type=pa.utf8()),
        "schema_path": pa.array([f"/tables/{n}/_schema" for n in names], type=pa.utf8()),
    })
    return arrow_response(table, negotiate(request))


@router.get("/{name}")
async def get_table(name: str, request: Request, columns: str | None = None, filters: str | None = None) -> Response:
    fn = request.app.state.data_sources.get(name)
    if fn is None:
        return Response(content=f"table not found: {name}", status_code=404)
    df = fn()
    if not isinstance(df, pd.DataFrame):
        raise TypeError(f"data source '{name}' must return a DataFrame, got {type(df).__name__}")
    table = pa.Table.from_pandas(df)
    table = apply_filters(table, filters)
    table = apply_columns(table, columns)
    return arrow_response(table, negotiate(request))


@router.get("/{name}/_schema")
async def get_schema(name: str, request: Request) -> Response:
    fn = request.app.state.data_sources.get(name)
    if fn is None:
        return Response(content=f"table not found: {name}", status_code=404)
    df = fn()
    if not isinstance(df, pd.DataFrame):
        raise TypeError(f"data source '{name}' must return a DataFrame, got {type(df).__name__}")
    table = pa.Table.from_pandas(df)
    empty = table.schema.empty_table()
    return arrow_response(empty, negotiate(request))
