tests/test_analysis_tools.py

from __future__ import annotations

from types import SimpleNamespace

import pytest
from pydantic import ValidationError

from backend.analysis_tools import (
    TOOL_BY_NAME,
    build_analysis,
    tool_declarations,
)
from backend.app import corrected_question, requested_city_names
from backend.database import AirQualityDatabase
from backend.gemini_service import GeminiService


@pytest.fixture(scope="module")
def database():
    instance = AirQualityDatabase()
    instance.initialize()
    return instance


@pytest.mark.parametrize(
    ("name", "arguments", "expected_columns"),
    [
        (
            "rank_cities",
            {
                "pollutant": "pm25",
                "start_year": 2023,
                "end_year": 2023,
                "limit": 10,
            },
            {
                "city",
                "mean_pm25",
                "station_count",
                "min_observation_days",
                "total_station_days",
            },
        ),
        (
            "city_average",
            {
                "pollutant": "pm25",
                "city": "Mumbai",
                "start_year": 2017,
                "end_year": 2024,
            },
            {
                "city",
                "mean_pm25",
                "station_count",
                "min_observation_days",
                "total_station_days",
            },
        ),
        (
            "threshold_cities",
            {
                "pollutant": "pm25",
                "threshold": 60,
                "start_year": 2023,
                "end_year": 2023,
            },
            {
                "city",
                "mean_pm25",
                "station_count",
                "min_observation_days",
                "matching_city_count",
            },
        ),
        (
            "compare_cities",
            {
                "pollutant": "pm25",
                "cities": ["delhi", "MUMBAI"],
                "start_year": 2023,
                "end_year": 2023,
            },
            {"city", "mean_pm25", "station_count", "min_observation_days"},
        ),
        (
            "time_trend",
            {
                "pollutant": "pm25",
                "city": "delhi",
                "start_year": 2023,
                "end_year": 2023,
                "interval": "monthly",
            },
            {"period", "mean_pm25", "station_count", "min_observation_days"},
        ),
        (
            "relationship",
            {
                "x_metric": "rainfall",
                "y_metric": "pm25",
                "city": "delhi",
                "months": [9, 6, 8, 7, 7],
            },
            {
                "city",
                "date",
                "rainfall",
                "pm25",
                "station_pair_count",
                "pearson_r",
                "paired_days",
            },
        ),
        (
            "strongest_weather_relationship",
            {"pollutant": "pm25", "city": "Delhi"},
            {
                "metric",
                "pearson_r",
                "paired_days",
                "city_count",
                "min_station_pairs",
            },
        ),
        (
            "seasonal_profile",
            {"pollutant": "pm25", "city": "Delhi"},
            {
                "season",
                "average_pm25",
                "city_count",
                "station_count",
                "min_observation_days",
            },
        ),
        (
            "funding_lookup",
            {"cities": ["Delhi", "Mumbai"]},
            {
                "city",
                "state",
                "total_fund_released",
                "utilisation_june_2022",
            },
        ),
        (
            "station_coverage",
            {
                "metric": "pm25",
                "cities": ["Mumbai", "Delhi"],
            },
            {
                "city",
                "station_count",
                "min_observation_days",
                "total_station_days",
            },
        ),
        (
            "weekday_weekend_profile",
            {"pollutant": "pm25", "city": "Delhi", "start_year": 2023, "end_year": 2023},
            {"day_type", "mean_pm25", "station_count", "city_count"},
        ),
        (
            "condition_comparison",
            {
                "pollutant": "pm25",
                "condition_metric": "wind_speed",
                "threshold": 3,
                "city": "Delhi",
            },
            {"condition_group", "mean_pm25", "station_count", "city_count"},
        ),
        (
            "rank_states",
            {"pollutant": "pm25", "start_year": 2023, "end_year": 2023},
            {"state", "mean_pm25", "city_count", "station_count"},
        ),
        (
            "coverage_trend",
            {"metric": "pm25"},
            {"year", "station_count", "city_count", "total_station_days"},
        ),
        (
            "funding_rank",
            {"limit": 10},
            {"city", "state", "total_fund_released"},
        ),
        (
            "ncap_threshold_cities",
            {"pollutant": "pm25", "threshold": 60},
            {"city", "mean_pm25", "total_fund_released", "station_count"},
        ),
        (
            "ncap_funding_groups",
            {"pollutant": "pm25"},
            {"funding_group", "mean_pm25", "city_count"},
        ),
        (
            "context_relationship",
            {
                "pollutant": "pm25",
                "context_metric": "total_fund_released",
                "pollution_measure": "level",
                "start_year": 2017,
                "end_year": 2024,
            },
            {"city", "total_fund_released", "mean_pm25", "pearson_r", "paired_cities"},
        ),
        (
            "pollution_change",
            {
                "pollutant": "pm25",
                "start_year": 2022,
                "end_year": 2023,
                "scope": "ncap_funded",
                "change_filter": "reductions_only",
            },
            {
                "city",
                "mean_pm25_2022",
                "mean_pm25_2023",
                "absolute_change_pm25",
                "percent_change",
                "matched_station_count",
                "total_fund_released",
            },
        ),
        (
            "context_relationship",
            {
                "pollutant": "pm25",
                "start_year": 2022,
                "end_year": 2023,
                "pollution_measure": "change",
                "context_metric": "total_fund_released",
            },
            {
                "city",
                "total_fund_released",
                "absolute_change_pm25",
                "pearson_r",
                "paired_cities",
            },
        ),
        (
            "threshold_frequency",
            {
                "pollutant": "pm25",
                "threshold": 60,
                "start_year": 2023,
                "end_year": 2023,
            },
            {
                "city",
                "threshold_day_count",
                "observed_day_count",
                "threshold_day_share_pct",
            },
        ),
        (
            "context_relationship",
            {
                "pollutant": "pm25",
                "context_metric": "population_density",
                "pollution_measure": "level",
                "start_year": 2023,
                "end_year": 2023,
            },
            {
                "state",
                "mean_pm25",
                "population_density",
                "pearson_r",
                "paired_states",
            },
        ),
    ],
)
def test_prebuilt_analysis_executes_safely(
    database,
    name,
    arguments,
    expected_columns,
):
    analysis = build_analysis(name, arguments)
    assert analysis.name == name
    safe_sql = database.validate_sql(analysis.plan.sql)
    columns, rows, truncated = database.execute(safe_sql)
    assert expected_columns.issubset(columns)
    assert rows
    assert truncated is False


def test_rank_cities_uses_equal_station_weighting_and_known_baseline(database):
    analysis = build_analysis(
        "rank_cities",
        {
            "pollutant": "pm25",
            "start_year": 2023,
            "end_year": 2023,
            "limit": 10,
        },
    )
    _, rows, _ = database.execute(analysis.plan.sql)
    assert rows[0] == {
        "city": "Byrnihat",
        "mean_pm25": 151.51,
        "station_count": 1,
        "min_observation_days": 351,
        "total_station_days": 351,
    }
    assert all(row["min_observation_days"] >= 30 for row in rows)


def test_single_city_average_and_station_coverage_have_known_baselines(database):
    average = build_analysis(
        "city_average",
        {"pollutant": "pm25", "city": "Mumbai"},
    )
    _, average_rows, _ = database.execute(average.plan.sql)
    assert len(average_rows) == 1
    assert average_rows[0]["city"] == "Mumbai"
    assert average_rows[0]["station_count"] == 30
    assert average_rows[0]["mean_pm25"] > 0

    coverage = build_analysis(
        "station_coverage",
        {"metric": "pm25", "cities": ["Mumbai", "Delhi"]},
    )
    _, coverage_rows, _ = database.execute(coverage.plan.sql)
    assert [(row["city"], row["station_count"]) for row in coverage_rows] == [
        ("Delhi", 38),
        ("Mumbai", 30),
    ]


def test_relationship_uses_all_pairs_for_statistic_before_chart_limit(database):
    analysis = build_analysis(
        "relationship",
        {
            "x_metric": "rainfall",
            "y_metric": "pm25",
            "months": [6, 7, 8, 9],
        },
    )
    _, rows, _ = database.execute(analysis.plan.sql)
    assert len(rows) == 100
    assert rows[0]["paired_days"] > len(rows)
    assert -1 <= rows[0]["pearson_r"] <= 1
    assert all(row["paired_days"] == rows[0]["paired_days"] for row in rows)


def test_strongest_weather_relationship_evaluates_all_factors(database):
    analysis = build_analysis(
        "strongest_weather_relationship",
        {"pollutant": "pm25", "start_year": 2023, "end_year": 2023},
    )
    _, rows, _ = database.execute(analysis.plan.sql)
    assert {row["metric"] for row in rows} == {
        "temperature",
        "humidity",
        "wind_speed",
        "rainfall",
        "solar_radiation",
        "pressure",
    }
    absolute_correlations = [abs(row["pearson_r"]) for row in rows]
    assert absolute_correlations == sorted(absolute_correlations, reverse=True)
    assert all(row["paired_days"] >= 30 for row in rows)


def test_ncap_change_uses_same_stations_and_reports_funding_context(database):
    analysis = build_analysis(
        "pollution_change",
        {
            "pollutant": "pm25",
            "start_year": 2022,
            "end_year": 2023,
            "scope": "ncap_funded",
            "change_filter": "reductions_only",
            "order": "largest_reduction",
            "limit": 10,
        },
    )
    _, rows, _ = database.execute(analysis.plan.sql)
    assert rows
    assert all(row["absolute_change_pm25"] < 0 for row in rows)
    assert all(row["matched_station_count"] >= 1 for row in rows)
    assert all(row["min_start_observation_days"] >= 180 for row in rows)
    assert all(row["min_end_observation_days"] >= 180 for row in rows)
    assert all(row["min_start_observation_months"] >= 9 for row in rows)
    assert all(row["min_end_observation_months"] >= 9 for row in rows)
    assert rows[0]["city"] == "Gaya"
    assert rows[0]["absolute_change_pm25"] == -11.37
    assert "only the same stations" in analysis.plan.method_note
    assert "does not attribute" in analysis.plan.method_note


def test_named_change_comparison_keeps_both_increase_and_reduction(database):
    analysis = build_analysis(
        "pollution_change",
        {
            "pollutant": "pm25",
            "start_year": 2022,
            "end_year": 2023,
            "cities": ["Delhi", "Mumbai"],
            "scope": "all_cities",
            "change_filter": "all",
            "order": "largest_reduction",
        },
    )
    _, rows, _ = database.execute(analysis.plan.sql)
    assert {row["city"] for row in rows} == {"Delhi", "Mumbai"}
    assert any(row["absolute_change_pm25"] < 0 for row in rows)
    assert any(row["absolute_change_pm25"] > 0 for row in rows)


def test_change_relationship_preserves_city_and_state_level_granularity(database):
    funding = build_analysis(
        "context_relationship",
        {
            "pollutant": "pm25",
            "start_year": 2022,
            "end_year": 2023,
            "pollution_measure": "change",
            "context_metric": "total_fund_released",
        },
    )
    _, funding_rows, _ = database.execute(funding.plan.sql)
    assert funding_rows[0]["paired_cities"] == len(funding_rows)
    assert all("city" in row for row in funding_rows)
    assert -1 <= funding_rows[0]["pearson_r"] <= 1

    utilisation = build_analysis(
        "context_relationship",
        {
            "pollutant": "pm25",
            "start_year": 2022,
            "end_year": 2023,
            "pollution_measure": "change",
            "context_metric": "utilisation_june_2022",
        },
    )
    _, utilisation_rows, _ = database.execute(utilisation.plan.sql)
    assert utilisation_rows[0]["paired_states"] == len(utilisation_rows)
    assert all("state" in row and "city" not in row for row in utilisation_rows)
    assert "each state was included once" in utilisation.plan.method_note


def test_exceedance_frequency_does_not_treat_missing_days_as_clean(database):
    analysis = build_analysis(
        "threshold_frequency",
        {
            "pollutant": "pm25",
            "threshold": 60,
            "start_year": 2023,
            "end_year": 2023,
            "rank_by": "share",
            "order": "most",
            "limit": 10,
        },
    )
    _, rows, _ = database.execute(analysis.plan.sql)
    assert rows
    assert all(row["observed_day_count"] >= 30 for row in rows)
    assert all(
        0 <= row["threshold_day_count"] <= row["observed_day_count"]
        for row in rows
    )
    assert all(0 <= row["threshold_day_share_pct"] <= 100 for row in rows)
    assert "Missing days were not treated" in analysis.plan.method_note


def test_utilisation_ranking_is_state_level_not_repeated_by_city(database):
    analysis = build_analysis(
        "funding_rank",
        {
            "metric": "utilisation_june_2022",
            "order": "highest",
            "limit": 10,
        },
    )
    _, rows, _ = database.execute(analysis.plan.sql)
    assert rows
    assert all("state" in row and "city" not in row for row in rows)
    assert len({row["state"] for row in rows}) == len(rows)
    assert "Ranked states, not cities" in analysis.plan.method_note


def test_city_typo_suggestions_come_from_real_dataset_names(database):
    assert database.suggest_city_names(["Dheli"]) == {"Dheli": "Delhi"}
    assert database.suggest_city_names(["Delhi"]) == {}
    assert database.suggest_city_names(["not a real place at all"]) == {}


def test_requested_city_extraction_and_corrected_follow_up():
    assert requested_city_names("time_trend", {"city": "Dheli"}) == ["Dheli"]
    assert requested_city_names(
        "compare_cities",
        {"cities": ["Dheli", "Mumbai"]},
    ) == ["Dheli", "Mumbai"]
    assert requested_city_names("rank_cities", {"limit": 10}) == []
    assert requested_city_names(
        "pollution_change",
        {"cities": ["Dheli"], "scope": "ncap_funded"},
    ) == ["Dheli"]
    assert corrected_question(
        "Show monthly PM2.5 for Dheli.",
        "Dheli",
        "Delhi",
    ) == "Show monthly PM2.5 for Delhi."


def test_threshold_reports_full_match_count_even_when_output_is_limited(database):
    analysis = build_analysis(
        "threshold_cities",
        {
            "pollutant": "pm25",
            "threshold": 20,
            "limit": 5,
        },
    )
    _, rows, _ = database.execute(analysis.plan.sql)
    assert len(rows) == 5
    assert rows[0]["matching_city_count"] > len(rows)


def test_city_values_are_sql_escaped_and_cannot_change_the_query(database):
    analysis = build_analysis(
        "time_trend",
        {
            "pollutant": "pm25",
            "city": "Delhi' OR 1=1",
            "start_year": 2023,
            "end_year": 2023,
        },
    )
    safe_sql = database.validate_sql(analysis.plan.sql)
    _, rows, _ = database.execute(safe_sql)
    assert rows == []


@pytest.mark.parametrize(
    ("name", "arguments"),
    [
        (
            "rank_cities",
            {"pollutant": "pm25", "start_year": 2024, "end_year": 2023},
        ),
        (
            "threshold_cities",
            {"pollutant": "pm25", "threshold": -1},
        ),
        (
            "compare_cities",
            {"pollutant": "pm25", "cities": ["Delhi", "delhi"]},
        ),
        (
            "relationship",
            {"x_metric": "pm25", "y_metric": "pm25"},
        ),
        (
            "time_trend",
            {"pollutant": "invalid", "city": "Delhi"},
        ),
    ],
)
def test_invalid_function_arguments_are_rejected(name, arguments):
    with pytest.raises(ValidationError):
        build_analysis(name, arguments)


def test_out_of_scope_can_return_a_helpful_explanation():
    definition = TOOL_BY_NAME["out_of_scope"]
    validated = definition.arguments_model.model_validate({"reason": "x" * 600})
    assert len(validated.reason) == 600

    with pytest.raises(ValidationError):
        definition.arguments_model.model_validate({"reason": "x" * 601})


def test_city_average_contract_makes_temporal_and_weighting_choices_explicit():
    analysis = build_analysis(
        "city_average",
        {
            "city": "Mumbai",
            "pollutant": "pm25",
            "start_year": 2020,
            "end_year": 2023,
            "months": [12, 1, 2],
            "statistic": "median",
            "minimum_station_days": 45,
            "station_weighting": "equal_station",
        },
    )
    assert "MEDIAN(pm25)" in analysis.plan.sql
    assert "year BETWEEN 2020 AND 2023" in analysis.plan.sql
    assert "IN (1, 2, 12)" in analysis.plan.sql
    assert ">= 45" in analysis.plan.sql
    assert "equal-station weighting" in analysis.plan.method_note
    assert "custom_sql_analysis" not in TOOL_BY_NAME


def test_tool_declarations_and_router_extraction_are_closed_over_known_tools():
    declarations = tool_declarations()
    assert {item["name"] for item in declarations} == set(TOOL_BY_NAME)
    assert all(item["type"] == "function" for item in declarations)
    assert all(item["parameters"]["additionalProperties"] is False for item in declarations)

    response = SimpleNamespace(
        steps=[
            SimpleNamespace(
                type="function_call",
                name="rank_cities",
                arguments={"pollutant": "pm25", "limit": 10},
            )
        ]
    )
    call = GeminiService._extract_tool_call(response)
    assert call.name == "rank_cities"
    assert call.arguments["limit"] == 10


def test_router_extraction_rejects_text_or_multiple_calls():
    with pytest.raises(ValueError, match="between one and three"):
        GeminiService._extract_tool_call(SimpleNamespace(steps=[]))
    with pytest.raises(ValueError, match="exactly one"):
        GeminiService._extract_tool_call(
            SimpleNamespace(
                steps=[
                    SimpleNamespace(
                        type="function_call",
                        name="rank_cities",
                        arguments={},
                    ),
                    SimpleNamespace(
                        type="function_call",
                        name="threshold_cities",
                        arguments={},
                    ),
                ]
            )
        )


def test_router_accepts_composition_of_up_to_three_typed_functions():
    response = SimpleNamespace(
        steps=[
            SimpleNamespace(
                type="function_call",
                name="city_average",
                arguments={"city": "Mumbai", "pollutant": "pm25"},
            ),
            SimpleNamespace(
                type="function_call",
                name="station_coverage",
                arguments={"cities": ["Mumbai"], "metric": "pm25"},
            ),
        ]
    )
    routed = GeminiService._extract_tool_calls(response)
    assert [call.name for call in routed.calls] == [
        "city_average",
        "station_coverage",
    ]
    assert routed.name == "composition"

    with pytest.raises(ValueError, match="cannot be composed"):
        GeminiService._extract_tool_calls(
            SimpleNamespace(
                steps=[
                    *response.steps,
                    SimpleNamespace(
                        type="function_call",
                        name="out_of_scope",
                        arguments={},
                    ),
                ]
            )
        )


def test_malformed_tool_call_errors_are_retryable_but_auth_errors_are_not():
    malformed = RuntimeError(
        "Model generated invalid JSON syntax: malformed_tool_call"
    )
    assert GeminiService._is_malformed_tool_call_error(malformed)
    assert not GeminiService._is_malformed_tool_call_error(
        RuntimeError("401 invalid API key")
    )