"""Data Connect v1 API client."""

from __future__ import annotations

from typing import Any

from ai4gh.core.client import GA4GHClient
from ai4gh.specs.data_connect.models import TableData, TableInfo, TableListResponse


class DataConnectClient(GA4GHClient):
    """Client for the GA4GH Data Connect v1 API.

    Data Connect provides a standard API for searching tabular data
    using SQL-like expressions across federated datasets.
    """

    spec_name = "data_connect"
    spec_version = "1.0.0"
    base_path = ""

    async def list_tables(self) -> TableListResponse:
        """List available tables.

        Returns:
            A list of table information objects.
        """
        data = await self._request("GET", "/tables")
        return TableListResponse.model_validate(data)

    async def get_table_info(self, table_name: str) -> TableInfo:
        """Get information about a specific table.

        Args:
            table_name: The table name.

        Returns:
            Table information including schema.
        """
        data = await self._request("GET", f"/table/{table_name}/info")
        return TableInfo.model_validate(data)

    async def get_table_data(
        self,
        table_name: str,
        *,
        page_size: int | None = None,
        page_token: str | None = None,
    ) -> TableData:
        """Get data from a table.

        Args:
            table_name: The table name.
            page_size: Maximum number of rows to return.
            page_token: Token for pagination.

        Returns:
            Table data with optional pagination.
        """
        params: dict[str, Any] = {}
        if page_size is not None:
            params["page_size"] = page_size
        if page_token is not None:
            params["page_token"] = page_token
        data = await self._request("GET", f"/table/{table_name}/data", params=params or None)
        return TableData.model_validate(data)

    async def search(
        self,
        query: str,
        *,
        page_size: int | None = None,
        page_token: str | None = None,
    ) -> TableData:
        """Search using a SQL-like expression.

        Args:
            query: SQL-like query string.
            page_size: Maximum number of rows to return.
            page_token: Token for pagination.

        Returns:
            Query results with optional pagination.
        """
        body: dict[str, Any] = {"query": query}
        if page_size is not None:
            body["page_size"] = page_size
        if page_token is not None:
            body["page_token"] = page_token
        data = await self._request("POST", "/search", json=body)
        return TableData.model_validate(data)

    async def get_next_page(self, pagination_url: str) -> TableData:
        """Fetch the next page of results.

        Args:
            pagination_url: The next_page_url from pagination metadata.

        Returns:
            The next page of data.
        """
        auth_headers = await self._get_auth_headers()
        response = await self._transport.get(pagination_url, headers=auth_headers)
        return TableData.model_validate(response.json())
