# tests/test_cendat.py

import pytest
import re
from unittest.mock import patch, Mock

from cendat.client import CenDatHelper, CenDatResponse

# --- 1. Mock Data ---

SIMPLE_PRODUCTS_JSON = {
    "dataset": [
        {
            "title": "American Community Survey",
            "c_isAggregate": "true",
            "c_vintage": 2022,
            "distribution": [{"accessURL": "http://api.census.gov/data/2022/acs/acs5"}],
        },
        {
            "title": "PUMS Household Data",
            "c_isMicrodata": "true",
            "c_vintage": 2022,
            "distribution": [
                {"accessURL": "http://api.census.gov/data/2022/acs/acs5/pums"}
            ],
        },
        {
            "title": "A Timeseries Dataset",  # This product should be filtered out
            "c_isAggregate": "false",
            "c_isMicrodata": "false",
            "c_vintage": 2022,
            "distribution": [
                {"accessURL": "http://api.census.gov/data/2022/timeseries"}
            ],
        },
    ]
}

FAKE_GEOS_JSON = {
    "fips": [
        {"name": "us", "geoLevelDisplay": "010", "requires": None},
        {"name": "state", "geoLevelDisplay": "040", "requires": None},
        {"name": "county", "geoLevelDisplay": "050", "requires": ["state"]},
        {"name": "tract", "geoLevelDisplay": "140", "requires": ["state", "county"]},
        {
            "name": "public use microdata area",
            "geoLevelDisplay": "795",
            "requires": ["state"],
        },
    ]
}

SIMPLE_VARIABLES_JSON = {
    "variables": {
        "B01001_001E": {"label": "Total Population", "concept": "SEX BY AGE"},
        "B19013_001E": {"label": "Median Household Income", "concept": "INCOME"},
        "PUMA": {"label": "Public Use Microdata Area Code", "concept": "GEOGRAPHY"},
    }
}


# --- 2. Pytest Fixtures ---
@pytest.fixture
def cdh():
    """Returns a fresh CenDatHelper instance for each test."""
    return CenDatHelper()


@pytest.fixture
def sample_response():
    """Returns a pre-populated CenDatResponse object for testing its methods."""
    data = [
        {
            "product": "Product A",
            "vintage": [2022],
            "sumlev": "040",
            "desc": "state",
            "schema": ["NAME", "POP"],
            "data": [["Alabama", "5000000"], ["Alaska", "700000"]],
        },
        {
            "product": "Product B",
            "vintage": [2022],
            "sumlev": "050",
            "desc": "county",
            "schema": ["NAME", "POP"],
            "data": [["Autauga County", "58000"]],
        },
    ]
    return CenDatResponse(data)


@pytest.fixture
def tabulation_response():
    """Returns a more complex CenDatResponse object for testing tabulation."""
    data = [
        {
            "product": "Test Survey",
            "vintage": [2022],
            "sumlev": "040",
            "desc": "state",
            "names": ["AGE", "WEIGHT"],  # For destringing
            "schema": ["STATE", "RACE", "AGE", "WEIGHT"],
            "data": [
                ["CA", "White", "25", "100"],
                ["CA", "Black", "35", "120"],
                ["TX", "White", "25", "80"],
                ["TX", "White", "45", "90"],
                ["FL", "Asian", "50", "110"],
            ],
        }
    ]
    return CenDatResponse(data)


# --- 3. Test Functions ---


@pytest.mark.unit
@patch("cendat.client.requests.get")
def test_list_products_filters_unsupported_types(mock_get, cdh):
    mock_response = Mock()
    mock_response.json.return_value = SIMPLE_PRODUCTS_JSON
    mock_get.return_value = mock_response
    products = cdh.list_products()
    assert len(products) == 2  # Should include aggregate and microdata
    assert products[0]["title"] == "American Community Survey (2022/acs/acs5)"
    assert products[1]["title"] == "PUMS Household Data (2022/acs/acs5/pums)"


@pytest.mark.unit
@patch("cendat.client.requests.get")
def test_list_variables_with_patterns_and_logic(mock_get, cdh):
    """Tests that the search and filter parameters work correctly."""
    mock_product_response = Mock()
    mock_product_response.json.return_value = SIMPLE_PRODUCTS_JSON
    mock_variable_response = Mock()
    mock_variable_response.json.return_value = SIMPLE_VARIABLES_JSON
    mock_get.side_effect = [mock_product_response, mock_variable_response]

    cdh.set_products(titles="American Community Survey (2022/acs/acs5)")
    variables = cdh.list_variables(
        patterns=["INCOME", "GEOGRAPHY"], logic=any, match_in="concept"
    )

    # FIX: Check for the presence of the correct names, as order is not guaranteed.
    returned_names = {v["name"] for v in variables}
    expected_names = {"B19013_001E", "PUMA"}

    assert len(variables) == 2
    assert returned_names == expected_names


@pytest.mark.unit
@patch("cendat.client.CenDatHelper._get_parent_geo_combinations")
@patch("cendat.client.CenDatHelper._get_json_from_url")  # Patch the internal helper
def test_get_data_preview_only_skips_fetching(mock_get_json, mock_get_combos, cdh):
    """Tests that preview_only=True sets n_calls but does not fetch data."""
    # This test is specifically for the preview logic, so we don't need ThreadPoolExecutor
    mock_get_json.side_effect = [
        SIMPLE_PRODUCTS_JSON,
        FAKE_GEOS_JSON,
        SIMPLE_VARIABLES_JSON,
    ]
    mock_get_combos.return_value = [{}]

    cdh.set_products(titles="American Community Survey (2022/acs/acs5)")
    cdh.set_geos(values="state", by="desc")
    cdh.set_variables(names="B01001_001E")

    # FIX: The method currently returns None in preview_only mode.
    # This test asserts the current behavior.
    response = cdh.get_data(preview_only=True)

    assert cdh.n_calls == 1
    # Check that the data-fetching part of _get_json_from_url was not called
    assert mock_get_json.call_count == 3  # Only setup calls
    assert response is None


@pytest.mark.unit
@patch("cendat.client.CenDatHelper._get_json_from_url")
def test_get_data_handles_microdata_correctly(mock_get_json, cdh):
    """Tests that the microdata workflow constructs the correct API call."""
    # Mock the setup calls first
    mock_get_json.side_effect = [
        SIMPLE_PRODUCTS_JSON,
        FAKE_GEOS_JSON,
        SIMPLE_VARIABLES_JSON,
        # Mock the two data calls with valid list-of-lists responses
        [["PUMA", "state"], ["01301", "08"]],
        [["PUMA", "state"], ["01302", "08"]],
    ]

    cdh.set_products(titles="PUMS Household Data (2022/acs/acs5/pums)")
    cdh.set_geos(values="public use microdata area", by="desc")
    cdh.set_variables(names="PUMA")
    response = cdh.get_data(
        within={"state": "08", "public use microdata area": ["01301", "01302"]}
    )

    assert mock_get_json.call_count == 5  # 3 setup calls + 2 data calls

    first_data_call = mock_get_json.call_args_list[-2]
    assert first_data_call.args[1]["for"] == "public use microdata area:01301"
    assert first_data_call.args[1]["in"] == "state:08"

    second_data_call = mock_get_json.call_args_list[-1]
    assert second_data_call.args[1]["for"] == "public use microdata area:01302"
    assert second_data_call.args[1]["in"] == "state:08"

    # Check that a valid response object was created
    assert isinstance(response, CenDatResponse)
    assert len(response._data[0]["data"]) == 2  # Both data rows should be aggregated


@pytest.mark.unit
@patch("cendat.client.requests.get")
def test_set_geos_requires_message(mock_get, cdh, capsys):
    mock_product_response = Mock()
    mock_product_response.json.return_value = {
        "dataset": [
            {
                "title": "American Community Survey",
                "c_isAggregate": "true",
                "c_vintage": 2022,
                "distribution": [
                    {"accessURL": "http://api.census.gov/data/2022/acs/acs5"}
                ],
            }
        ]
    }
    mock_product_response.raise_for_status.return_value = None
    mock_geo_response = Mock()
    mock_geo_response.json.return_value = FAKE_GEOS_JSON
    mock_geo_response.raise_for_status.return_value = None
    mock_get.side_effect = [mock_product_response, mock_geo_response]

    cdh.set_products(titles="American Community Survey (2022/acs/acs5)")
    cdh.set_geos(values="tract", by="desc")

    captured = capsys.readouterr()
    expected_message = (
        "✅ Geographies set: 'tract' (requires `within` for: county, state)\n"
    )
    assert captured.out.endswith(expected_message)


@pytest.mark.unit
@patch("cendat.client.CenDatHelper._get_parent_geo_combinations")
@patch("cendat.client.CenDatHelper._get_json_from_url")
def test_get_data_expands_within_clauses_correctly(mock_get_json, mock_get_combos, cdh):
    mock_get_json.side_effect = [
        SIMPLE_PRODUCTS_JSON,
        FAKE_GEOS_JSON,
        SIMPLE_VARIABLES_JSON,
        # Mock data responses for the two API calls
        [["B01001_001E"], ["100"]],
        [["B01001_001E"], ["200"]],
    ]
    mock_get_combos.return_value = [{"state": "08", "county": "123"}]

    cdh.set_products(titles="American Community Survey (2022/acs/acs5)")
    cdh.set_geos(values="tract", by="desc")
    cdh.set_variables(names="B01001_001E")
    response = cdh.get_data(within=[{"state": "08", "county": "123"}, {"state": "56"}])

    assert mock_get_combos.call_count == 2
    first_call_args = mock_get_combos.call_args_list[0]
    assert first_call_args.args[2] == {"state": "08", "county": "123"}
    second_call_args = mock_get_combos.call_args_list[1]
    assert second_call_args.args[2] == {"state": "56"}
    assert isinstance(response, CenDatResponse)
    assert len(response._data[0]["data"]) == 2


@pytest.mark.unit
@patch("cendat.client.CenDatHelper._get_parent_geo_combinations")
@patch("cendat.client.CenDatHelper._get_json_from_url")
def test_get_data_handles_complex_list_expansion(mock_get_json, mock_get_combos, cdh):
    mock_get_json.side_effect = [
        SIMPLE_PRODUCTS_JSON,
        FAKE_GEOS_JSON,
        SIMPLE_VARIABLES_JSON,
        # Mock 4 data responses
        [["B01001_001E"], ["1"]],
        [["B01001_001E"], ["2"]],
        [["B01001_001E"], ["3"]],
        [["B01001_001E"], ["4"]],
    ]
    mock_get_combos.return_value = [{}]

    cdh.set_products(titles="American Community Survey (2022/acs/acs5)")
    cdh.set_geos(values="tract", by="desc")
    cdh.set_variables(names="B01001_001E")
    response = cdh.get_data(
        within=[{"state": "08", "county": ["069", "123"]}, {"state": ["36", "06"]}]
    )

    assert mock_get_combos.call_count == 4
    call_args = [call.args for call in mock_get_combos.call_args_list]
    assert call_args[0][2] == {"state": "08", "county": "069"}
    assert call_args[1][2] == {"state": "08", "county": "123"}
    assert call_args[2][2] == {"state": "36"}
    assert call_args[3][2] == {"state": "06"}
    assert isinstance(response, CenDatResponse)
    assert len(response._data[0]["data"]) == 4


@pytest.mark.unit
def test_to_pandas_concatenation(sample_response):
    """Tests that concat=True returns a single pandas DataFrame."""
    import pandas as pd

    df = sample_response.to_pandas(concat=True)
    assert isinstance(df, pd.DataFrame)
    assert len(df) == 3


@pytest.mark.unit
def test_to_polars_schema_overrides(sample_response):
    """Tests that schema_overrides correctly casts types in Polars."""
    import polars as pl

    df = sample_response.to_polars(schema_overrides={"POP": pl.Int64}, concat=True)
    assert df["POP"].dtype == pl.Int64


# --- 4. New Tests for Tabulation and Where Clauses ---


def find_row_in_output(pattern, text):
    """Helper to find a specific row in the formatted table output."""
    # Find lines that look like table rows │ ... │
    rows = re.findall(r"│(.*?)│", text)
    for row in rows:
        if re.search(pattern, row):
            return True
    return False


@pytest.mark.unit
def test_tabulate_simple_unweighted(tabulation_response, capsys):
    """Tests a simple tabulation without weights."""
    tabulation_response.tabulate("STATE", "RACE")
    captured = capsys.readouterr().out
    assert find_row_in_output(r"CA\s+┆\s+Black\s+┆\s+1", captured)
    assert find_row_in_output(r"CA\s+┆\s+White\s+┆\s+1", captured)
    assert find_row_in_output(r"TX\s+┆\s+White\s+┆\s+2", captured)
    assert find_row_in_output(r"FL\s+┆\s+Asian\s+┆\s+1", captured)


@pytest.mark.unit
def test_tabulate_weighted(tabulation_response, capsys):
    """Tests a tabulation with a weight variable."""
    tabulation_response.tabulate("STATE", weight_var="WEIGHT")
    captured = capsys.readouterr().out
    assert find_row_in_output(r"CA\s+┆\s+220", captured)
    assert find_row_in_output(r"TX\s+┆\s+170", captured)
    assert find_row_in_output(r"FL\s+┆\s+110", captured)


@pytest.mark.unit
def test_tabulate_weighted_with_divisor(tabulation_response, capsys):
    """Tests a tabulation with a weight variable and a divisor."""
    tabulation_response.tabulate("STATE", weight_var="WEIGHT", weight_div=10)
    captured = capsys.readouterr().out
    assert find_row_in_output(r"CA\s+┆\s+22.0", captured)
    assert find_row_in_output(r"TX\s+┆\s+17.0", captured)
    assert find_row_in_output(r"FL\s+┆\s+11.0", captured)


@pytest.mark.unit
def test_tabulate_with_where_clause(tabulation_response, capsys):
    """Tests that the where clause correctly filters data before tabulation."""
    tabulation_response.tabulate("STATE", where="AGE > 30")
    captured = capsys.readouterr().out
    assert find_row_in_output(r"CA\s+┆\s+1", captured)
    assert find_row_in_output(r"TX\s+┆\s+1", captured)
    assert find_row_in_output(r"FL\s+┆\s+1", captured)
    assert "White" not in captured  # The grouping variable is STATE, not RACE


@pytest.mark.unit
def test_where_clause_parser_valid(tabulation_response):
    """Tests that the _build_safe_checker method correctly parses valid conditions."""
    row1 = {"STATE": "CA", "RACE": "White", "AGE": 25, "WEIGHT": 100}
    row2 = {"STATE": "TX", "RACE": "Black", "AGE": 40, "WEIGHT": 120}

    checker_gt = tabulation_response._build_safe_checker("AGE > 30")
    assert not checker_gt(row1)
    assert checker_gt(row2)

    checker_eq = tabulation_response._build_safe_checker("RACE == 'White'")
    assert checker_eq(row1)
    assert not checker_eq(row2)

    checker_in = tabulation_response._build_safe_checker("STATE in ['CA', 'FL']")
    assert checker_in(row1)
    assert not checker_in(row2)


@pytest.mark.unit
def test_where_clause_parser_invalid(tabulation_response):
    """Tests that _build_safe_checker raises ValueError for invalid conditions."""
    # This should fail because the column doesn't exist, leading to a format error
    with pytest.raises(ValueError, match="Invalid condition format"):
        tabulation_response._build_safe_checker("INVALID_COL > 30")

    with pytest.raises(ValueError, match="Invalid condition format"):
        tabulation_response._build_safe_checker("AGE greater than 30")

    with pytest.raises(ValueError, match="Invalid value format"):
        tabulation_response._build_safe_checker("AGE > some_undefined_variable")


@pytest.mark.unit
def test_tabulate_with_multiple_where_clauses_and_any_logic(
    tabulation_response, capsys
):
    """Tests tabulation with a list of where clauses and 'any' logic."""
    tabulation_response.tabulate(
        "STATE", "RACE", where=["AGE < 30", "RACE == 'Asian'"], logic=any
    )
    captured = capsys.readouterr().out
    # Should include two people with AGE < 30 and one person who is Asian
    assert find_row_in_output(r"CA\s+┆\s+White\s+┆\s+1", captured)
    assert find_row_in_output(r"TX\s+┆\s+White\s+┆\s+1", captured)
    assert find_row_in_output(r"FL\s+┆\s+Asian\s+┆\s+1", captured)
    # The Black person in CA (age 35) should be excluded
    assert not find_row_in_output(r"CA\s+┆\s+Black", captured)


@pytest.mark.unit
@patch("cendat.client.CenDatHelper._get_parent_geo_combinations")
@patch("cendat.client.requests.get")
def test_get_data_with_specific_target_geos(mock_get, mock_get_combos, cdh):
    """
    Tests that get_data bypasses the parent combination search when the
    target geography is specified directly in the `within` clause.
    """
    # --- Arrange ---
    mock_product_response = Mock()
    mock_product_response.json.return_value = {
        "dataset": [
            {
                "title": "American Community Survey",
                "c_isAggregate": "true",
                "c_vintage": 2022,
                "distribution": [
                    {"accessURL": "http://api.census.gov/data/2022/acs/acs5"}
                ],
            }
        ]
    }
    mock_geo_response = Mock()
    mock_geo_response.json.return_value = FAKE_GEOS_JSON
    mock_variable_response = Mock()
    mock_variable_response.json.return_value = SIMPLE_VARIABLES_JSON

    mock_data_response = Mock()
    mock_data_response.json.return_value = [
        ["B01001_001E", "state", "county", "tract"],
        ["1234", "08", "069", "001201"],
    ]

    mock_get.side_effect = [
        mock_product_response,
        mock_geo_response,
        mock_variable_response,
        mock_data_response,
    ]

    # --- Act ---
    cdh.set_products(titles="American Community Survey (2022/acs/acs5)")
    cdh.set_geos(values="tract", by="desc")
    cdh.set_variables(names="B01001_001E")
    cdh.get_data(within={"state": "08", "county": "069", "tract": "001201"})

    # --- Assert ---
    mock_get_combos.assert_not_called()

    final_call_args = mock_get.call_args
    assert final_call_args.kwargs["params"]["for"] == "tract:001201"
    assert final_call_args.kwargs["params"]["in"] == "state:08 county:069"
