"""Tests for Data Connect client."""

from __future__ import annotations

import pytest
from pytest_httpx import HTTPXMock

from ai4gh.specs.data_connect.client import DataConnectClient


@pytest.fixture
def dc_client() -> DataConnectClient:
    return DataConnectClient(base_url="https://example.com")


TABLE_LIST_RESPONSE = {
    "tables": [
        {"name": "patients", "description": "Patient records"},
        {"name": "variants", "description": "Genomic variants"},
    ],
}

TABLE_INFO_RESPONSE = {
    "name": "patients",
    "description": "Patient records",
    "data_model": {
        "columns": [
            {"name": "id", "data_type": "string"},
            {"name": "age", "data_type": "integer"},
        ]
    },
}

TABLE_DATA_RESPONSE = {
    "data": [
        {"id": "1", "name": "Alice", "age": 30},
        {"id": "2", "name": "Bob", "age": 25},
    ],
    "pagination": {"next_page_url": "https://example.com/table/patients/data?page=2"},
}


class TestDataConnectClient:
    def test_spec_attrs(self) -> None:
        client = DataConnectClient(base_url="https://example.com")
        assert client.spec_name == "data_connect"
        assert client.spec_version == "1.0.0"
        assert client.base_path == ""

    def test_url(self) -> None:
        client = DataConnectClient(base_url="https://example.com")
        assert client.url == "https://example.com"

    async def test_list_tables(self, dc_client: DataConnectClient, httpx_mock: HTTPXMock) -> None:
        httpx_mock.add_response(
            url="https://example.com/tables",
            json=TABLE_LIST_RESPONSE,
        )
        resp = await dc_client.list_tables()
        assert len(resp.tables) == 2
        assert resp.tables[0].name == "patients"

    async def test_get_table_info(
        self, dc_client: DataConnectClient, httpx_mock: HTTPXMock
    ) -> None:
        httpx_mock.add_response(
            url="https://example.com/table/patients/info",
            json=TABLE_INFO_RESPONSE,
        )
        info = await dc_client.get_table_info("patients")
        assert info.name == "patients"
        assert len(info.data_model.columns) == 2

    async def test_get_table_data(
        self, dc_client: DataConnectClient, httpx_mock: HTTPXMock
    ) -> None:
        httpx_mock.add_response(
            url="https://example.com/table/patients/data",
            json=TABLE_DATA_RESPONSE,
        )
        data = await dc_client.get_table_data("patients")
        assert len(data.data) == 2
        assert data.data[0]["name"] == "Alice"

    async def test_get_table_data_with_pagination(
        self, dc_client: DataConnectClient, httpx_mock: HTTPXMock
    ) -> None:
        httpx_mock.add_response(json=TABLE_DATA_RESPONSE)
        await dc_client.get_table_data("patients", page_size=10, page_token="tok-1")
        request = httpx_mock.get_requests()[0]
        url = str(request.url)
        assert "page_size=10" in url
        assert "page_token=tok-1" in url

    async def test_search(self, dc_client: DataConnectClient, httpx_mock: HTTPXMock) -> None:
        httpx_mock.add_response(
            url="https://example.com/search",
            method="POST",
            json=TABLE_DATA_RESPONSE,
        )
        data = await dc_client.search("SELECT * FROM patients WHERE age > 25")
        assert len(data.data) == 2
        request = httpx_mock.get_requests()[0]
        import json

        body = json.loads(request.content)
        assert body["query"] == "SELECT * FROM patients WHERE age > 25"

    async def test_search_with_pagination(
        self, dc_client: DataConnectClient, httpx_mock: HTTPXMock
    ) -> None:
        httpx_mock.add_response(json=TABLE_DATA_RESPONSE)
        await dc_client.search(
            "SELECT * FROM patients",
            page_size=50,
            page_token="page-2",
        )
        request = httpx_mock.get_requests()[0]
        import json

        body = json.loads(request.content)
        assert body["page_size"] == 50
        assert body["page_token"] == "page-2"

    async def test_get_next_page(self, dc_client: DataConnectClient, httpx_mock: HTTPXMock) -> None:
        httpx_mock.add_response(
            url="https://example.com/table/patients/data?page=2",
            json={
                "data": [{"id": "3", "name": "Charlie"}],
            },
        )
        data = await dc_client.get_next_page("https://example.com/table/patients/data?page=2")
        assert len(data.data) == 1
        assert data.data[0]["name"] == "Charlie"
