"""Tests for aablocks.client module."""

from __future__ import annotations

import gzip
import json
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest
import requests

from aablocks.client import (
    DataBlock,
    download_structure,
    get_data,
    get_datacard,
    get_details,
    list_blocks,
)

from .conftest import SAMPLE_BLOCK_DICT, SAMPLE_BLOCKS_RESPONSE


# ---------------------------------------------------------------------------
# DataBlock dataclass
# ---------------------------------------------------------------------------
class TestDataBlock:
    def test_from_dict(self):
        """Constructs a DataBlock with all fields from a complete API response dict."""
        ds = DataBlock.from_dict(SAMPLE_BLOCK_DICT)
        assert ds.id == "ab1001"
        assert ds.name == "Test Data Block"
        assert ds.experiment == "Yeast Display"
        assert ds.modes == ["source", "ml"]
        assert ds.version == "3"
        assert ds.locked is False
        assert ds.url == "https://data.example.com/ab1001"
        assert ds.structure_count == 10000
        assert ds.a_size == 500
        assert ds.alpha_size == 200
        assert ds.total_ppi_count == 100000
        assert ds.unique_ppi_count == 50000
        assert ds.product == "atlas-vhh"
        assert ds.product_display_name == "Atlas VHH Consortia"
        assert ds.product_kind == "consortium"
        assert ds.source == "VHH Q1 2025"

    def test_from_dict_optional_fields(self):
        """Optional fields default to None/False when missing from the dict."""
        minimal = {
            "id": "ab9999",
            "name": "Minimal",
            "experiment": "",
            "details": "",
            "modes": [],
            "version": "1",
        }
        ds = DataBlock.from_dict(minimal)
        assert ds.locked is False
        assert ds.url is None
        assert ds.structure_count is None
        assert ds.a_size is None
        assert ds.alpha_size is None
        assert ds.total_ppi_count is None
        assert ds.unique_ppi_count is None
        assert ds.product is None
        assert ds.product_display_name is None
        assert ds.product_kind is None
        assert ds.source is None


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _json_response(data, status_code=200):
    resp = MagicMock(spec=requests.Response)
    resp.status_code = status_code
    resp.ok = status_code < 400
    resp.json.return_value = data
    resp.text = json.dumps(data)
    resp.raise_for_status.return_value = None
    return resp


# ---------------------------------------------------------------------------
# list_blocks
# ---------------------------------------------------------------------------
class TestListDataBlocks:
    def test_invalid_format_raises(self):
        """Passing an unsupported format raises ValueError before any API call."""
        with pytest.raises(ValueError, match="format must be"):
            list_blocks(format="xml")

    @patch("aablocks.client.requests.get")
    def test_list_format(self, mock_get):
        """format='list' returns a list of DataBlock objects."""
        mock_get.return_value = _json_response(SAMPLE_BLOCKS_RESPONSE)
        result = list_blocks(format="list")
        assert len(result) == 2
        assert isinstance(result[0], DataBlock)
        assert result[0].id == "ab1001"
        assert result[1].id == "ab1002"

    @patch("aablocks.client.requests.get")
    def test_csv_format(self, mock_get):
        """format='csv' returns a CSV string containing all block IDs."""
        mock_get.return_value = _json_response(SAMPLE_BLOCKS_RESPONSE)
        result = list_blocks(format="csv")
        assert isinstance(result, str)
        assert "ab1001" in result
        assert "ab1002" in result

    @patch("aablocks.client.requests.get")
    def test_pandas_format(self, mock_get):
        """format='pandas' returns a pandas DataFrame with correct columns."""
        pd = pytest.importorskip("pandas")
        mock_get.return_value = _json_response(SAMPLE_BLOCKS_RESPONSE)
        result = list_blocks(format="pandas")
        assert isinstance(result, pd.DataFrame)
        assert len(result) == 2
        assert list(result["id"]) == ["ab1001", "ab1002"]

    @patch("aablocks.client.requests.get")
    def test_polars_format(self, mock_get):
        """format='polars' returns a polars DataFrame with the right row count."""
        pl = pytest.importorskip("polars")
        mock_get.return_value = _json_response(SAMPLE_BLOCKS_RESPONSE)
        result = list_blocks(format="polars")
        assert isinstance(result, pl.DataFrame)
        assert len(result) == 2

    @patch("aablocks.client.requests.get")
    def test_all_versions_param(self, mock_get):
        """all_versions=True sends all_versions=true as a query parameter."""
        mock_get.return_value = _json_response({"objects": []})
        list_blocks(all_versions=True, format="list")
        _, kwargs = mock_get.call_args
        assert kwargs["params"] == {"all_versions": "true"}

    @patch("aablocks.client.requests.get")
    def test_empty_response(self, mock_get):
        """An empty objects array returns an empty list."""
        mock_get.return_value = _json_response({"objects": []})
        result = list_blocks(format="list")
        assert result == []

    @patch("aablocks.client.requests.get")
    def test_http_error_propagates(self, mock_get):
        """HTTP errors from the API are re-raised to the caller."""
        resp = MagicMock(spec=requests.Response)
        resp.status_code = 401
        resp.raise_for_status.side_effect = requests.HTTPError(response=resp)
        mock_get.return_value = resp
        with pytest.raises(requests.HTTPError):
            list_blocks(format="list")


# ---------------------------------------------------------------------------
# get_details
# ---------------------------------------------------------------------------
class TestGetDetails:
    @patch("aablocks.client.requests.get")
    def test_returns_data_block(self, mock_get):
        """Returns a DataBlock object populated from the API response."""
        mock_get.return_value = _json_response({"dataset": SAMPLE_BLOCK_DICT})
        ds = get_details("ab1001")
        assert ds.id == "ab1001"
        assert ds.name == "Test Data Block"

    @patch("aablocks.client.requests.get")
    def test_with_version_param(self, mock_get):
        """Passing a version sends it as a query parameter."""
        mock_get.return_value = _json_response({"dataset": SAMPLE_BLOCK_DICT})
        get_details("ab1001", version="2")
        _, kwargs = mock_get.call_args
        assert kwargs["params"] == {"version": "2"}

    @patch("aablocks.client.requests.get")
    def test_404_raises(self, mock_get):
        """A 404 response raises HTTPError for a nonexistent data block."""
        resp = MagicMock(spec=requests.Response)
        resp.status_code = 404
        resp.raise_for_status.side_effect = requests.HTTPError(response=resp)
        mock_get.return_value = resp
        with pytest.raises(requests.HTTPError):
            get_details("nonexistent")


# ---------------------------------------------------------------------------
# get_datacard
# ---------------------------------------------------------------------------
class TestGetDatacard:
    @patch("aablocks.client.requests.get")
    def test_returns_dict(self, mock_get):
        """Returns the parsed JSON body from the API response."""
        payload = {
            "exec": {
                "scientific_value": "VP text",
                "use_cases": [{"title": "UC1", "description": "desc"}],
            },
        }
        mock_get.return_value = _json_response(payload)
        result = get_datacard("ab1001")
        assert result == payload

    @patch("aablocks.client.requests.get")
    def test_with_version_param(self, mock_get):
        """Passing a version sends it as a query parameter."""
        mock_get.return_value = _json_response({"exec": {"scientific_value": "v1"}})
        get_datacard("ab1001", version="1")
        _, kwargs = mock_get.call_args
        assert kwargs["params"] == {"version": "1"}


# ---------------------------------------------------------------------------
# get_data
# ---------------------------------------------------------------------------
class TestGetDataBlock:
    def test_invalid_format_raises(self):
        """Passing an unsupported format raises ValueError before any API call."""
        with pytest.raises(ValueError, match="format must be"):
            get_data("ab1001", format="xml")

    @patch("aablocks.client.requests.get")
    def test_csv_s3_download(self, mock_get):
        """The API returns a presigned S3 URL; client downloads and decompresses."""
        csv_data = "col1,col2\nx,y\n"
        mock_get.side_effect = _mock_s3_download(csv_data)

        result = get_data("ab1001", format="csv", progress=False, load_negatives=True)
        assert result == csv_data

    @patch("aablocks.client.requests.get")
    def test_output_path_compressed(self, mock_get, tmp_path: Path):
        """Writing to file with output_compressed=True keeps the gzip encoding."""
        csv_data = "col1,col2\nx,y\n"
        mock_get.side_effect = _mock_s3_download(csv_data)

        out = tmp_path / "data.csv.gz"
        result = get_data(
            "ab1001",
            output_path=str(out),
            output_compressed=True,
            progress=False,
            load_negatives=True,
        )
        assert result is None
        assert gzip.decompress(out.read_bytes()).decode() == csv_data

    @patch("aablocks.client.requests.get")
    def test_output_path_decompressed(self, mock_get, tmp_path: Path):
        """Writing to file with output_compressed=False decompresses before writing."""
        csv_data = "col1,col2\nx,y\n"
        mock_get.side_effect = _mock_s3_download(csv_data)

        out = tmp_path / "data.csv"
        result = get_data(
            "ab1001",
            output_path=str(out),
            output_compressed=False,
            progress=False,
            load_negatives=True,
        )
        assert result is None
        assert out.read_text() == csv_data

    @patch("aablocks.client.requests.get")
    def test_pandas_s3_download(self, mock_get):
        """format='pandas' downloads from S3 and returns a DataFrame."""
        pd = pytest.importorskip("pandas")
        csv_data = "col1,col2\na,b\nc,d\n"
        mock_get.side_effect = _mock_s3_download(csv_data)

        # progress=True routes through _download_bytes (requests.get) so the
        # mocked S3 response is honored; progress=False would call pd.read_csv
        # on the URL directly and actually hit the network.
        result = get_data("ab1001", format="pandas", progress=True, load_negatives=True)
        assert isinstance(result, pd.DataFrame)
        assert list(result.columns) == ["col1", "col2"]
        assert len(result) == 2

    @patch("aablocks.client.requests.get")
    def test_params_sent_correctly(self, mock_get):
        """Version and mode are forwarded as query parameters on the API call."""
        mock_get.side_effect = _mock_s3_download("col1\na\n")

        get_data(
            "ab1001",
            version="2",
            mode="ml",
            format="csv",
            progress=False,
            load_negatives=True,
        )
        # The first call is the API; the second is the S3 download. Only the
        # first carries our params.
        api_call_kwargs = mock_get.call_args_list[0].kwargs
        assert api_call_kwargs["params"] == {
            "redirect": "false",
            "version": "2",
            "mode": "ml",
        }


# ---------------------------------------------------------------------------
# load_negatives filtering
# ---------------------------------------------------------------------------
CSV_WITH_NEGATIVES = (
    "mata_description,matalpha_description,value\n"
    "ProteinA,TargetX,1.0\n"
    "ANeg1,TargetY,0.1\n"
    "ProteinB,AlphaNeg1,0.2\n"
    "ANeg2,AlphaNeg2,0.05\n"
    "ProteinC,TargetZ,2.0\n"
)

CSV_WITHOUT_NEGATIVES = (
    "mata_description,matalpha_description,value\n"
    "ProteinA,TargetX,1.0\n"
    "ProteinC,TargetZ,2.0\n"
)


def _mock_s3_download(csv_text):
    """Return mock responses for the API call + S3 download flow."""
    gz_data = gzip.compress(csv_text.encode())

    api_resp = _json_response(
        {"url": "https://s3.example.com/data.csv.gz", "schema": {}}
    )

    s3_resp = MagicMock()
    s3_resp.__enter__ = MagicMock(return_value=s3_resp)
    s3_resp.__exit__ = MagicMock(return_value=False)
    s3_resp.raise_for_status.return_value = None
    s3_resp.headers = {"Content-Length": str(len(gz_data))}
    s3_resp.iter_content.return_value = [gz_data]

    return [api_resp, s3_resp]


class TestLoadNegatives:
    @patch("aablocks.client.requests.get")
    def test_csv_filters_negatives_by_default(self, mock_get):
        """load_negatives=False (default) removes ANeg/AlphaNeg rows from CSV output."""
        mock_get.side_effect = _mock_s3_download(CSV_WITH_NEGATIVES)
        result = get_data("ab1001", format="csv", progress=False, load_negatives=False)
        assert result == CSV_WITHOUT_NEGATIVES

    @patch("aablocks.client.requests.get")
    def test_csv_keeps_negatives_when_requested(self, mock_get):
        """load_negatives=True keeps all rows including negative controls."""
        mock_get.side_effect = _mock_s3_download(CSV_WITH_NEGATIVES)
        result = get_data("ab1001", format="csv", progress=False, load_negatives=True)
        assert result == CSV_WITH_NEGATIVES

    @patch("aablocks.client.requests.get")
    def test_pandas_filters_negatives(self, mock_get):
        """load_negatives=False removes ANeg/AlphaNeg rows from pandas DataFrame."""
        pd = pytest.importorskip("pandas")
        mock_get.side_effect = _mock_s3_download(CSV_WITH_NEGATIVES)
        result = get_data(
            "ab1001", format="pandas", progress=True, load_negatives=False
        )
        assert isinstance(result, pd.DataFrame)
        assert len(result) == 2
        assert list(result["mata_description"]) == ["ProteinA", "ProteinC"]

    @patch("aablocks.client.requests.get")
    def test_polars_filters_negatives(self, mock_get):
        """load_negatives=False removes ANeg/AlphaNeg rows from polars DataFrame."""
        pl = pytest.importorskip("polars")
        mock_get.side_effect = _mock_s3_download(CSV_WITH_NEGATIVES)
        result = get_data(
            "ab1001", format="polars", progress=False, load_negatives=False
        )
        assert isinstance(result, pl.DataFrame)
        assert len(result) == 2
        assert result["mata_description"].to_list() == ["ProteinA", "ProteinC"]

    @patch("aablocks.client.requests.get")
    def test_missing_columns_raises(self, mock_get):
        """Filtering raises ValueError when the data block lacks the expected columns."""
        csv_no_cols = "col1,col2\na,b\nc,d\n"
        mock_get.side_effect = _mock_s3_download(csv_no_cols)
        with pytest.raises(ValueError, match="missing column"):
            get_data("ab1001", format="csv", progress=False, load_negatives=False)


# ---------------------------------------------------------------------------
# download_structure
# ---------------------------------------------------------------------------
class TestDownloadStructure:
    @patch("aablocks.client.requests.get")
    def test_returns_content_without_output_path(self, mock_get):
        """Without output_path, returns the file content as a string."""
        resp = MagicMock(spec=requests.Response)
        resp.status_code = 200
        resp.text = "HEADER structure data\nATOM 1 CA\n"
        resp.raise_for_status.return_value = None
        mock_get.return_value = resp

        result = download_structure("ab1614", "structure.cif")
        assert result == resp.text

    @patch("aablocks.client.requests.get")
    def test_downloads_to_output_path(self, mock_get, tmp_path: Path):
        """output_path overrides the default filename."""
        resp = MagicMock(spec=requests.Response)
        resp.status_code = 200
        resp.text = "HEADER data\n"
        resp.raise_for_status.return_value = None
        mock_get.return_value = resp

        out = tmp_path / "custom.cif"
        result = download_structure("ab1614", "structure.cif", output_path=str(out))
        assert result == str(out)
        assert out.read_text() == resp.text

    @patch("aablocks.client.requests.get")
    def test_version_param(self, mock_get, tmp_path: Path):
        """Passing a version sends it as a query parameter."""
        resp = MagicMock(spec=requests.Response)
        resp.status_code = 200
        resp.text = "data"
        resp.raise_for_status.return_value = None
        mock_get.return_value = resp

        out = tmp_path / "s.cif"
        download_structure("ab1614", "s.cif", version="2", output_path=str(out))
        _, kwargs = mock_get.call_args
        assert kwargs["params"] == {"version": "2"}
