import polars as pl
import pytest

from hstrat._auxiliary_lib import alifestd_check_topological_sensitivity_polars
from hstrat._auxiliary_lib._alifestd_check_topological_sensitivity import (
    _topologically_sensitive_cols,
    _update_only_sensitive_cols,
)


@pytest.fixture
def base_df():
    return pl.DataFrame(
        {
            "id": [0, 1, 2],
            "ancestor_id": [0, 0, 1],
            "origin_time": [0, 1, 2],
        }
    )


def test_none_present(base_df: pl.DataFrame):
    result = alifestd_check_topological_sensitivity_polars(
        base_df,
        insert=True,
        delete=True,
        update=True,
    )
    assert result == []


@pytest.mark.parametrize("col", sorted(_topologically_sensitive_cols))
def test_single_sensitive_col(base_df: pl.DataFrame, col: str):
    df = base_df.with_columns(pl.lit(0).alias(col))
    result = alifestd_check_topological_sensitivity_polars(
        df,
        insert=True,
        delete=True,
        update=True,
    )
    assert result == [col]


def test_multiple_sensitive_cols(base_df: pl.DataFrame):
    df = base_df.with_columns(
        pl.lit(0).alias("branch_length"),
        pl.lit(0).alias("node_depth"),
        pl.lit(0).alias("sister_id"),
    )
    result = alifestd_check_topological_sensitivity_polars(
        df,
        insert=True,
        delete=True,
        update=True,
    )
    assert set(result) == {"branch_length", "node_depth", "sister_id"}


def test_non_sensitive_cols_ignored(base_df: pl.DataFrame):
    df = base_df.with_columns(
        pl.lit("x").alias("taxon_label"),
        pl.lit(True).alias("extant"),
    )
    result = alifestd_check_topological_sensitivity_polars(
        df,
        insert=True,
        delete=True,
        update=True,
    )
    assert result == []


def test_no_mutation(base_df: pl.DataFrame):
    df = base_df.with_columns(pl.lit(0).alias("branch_length"))
    original = df.clone()
    alifestd_check_topological_sensitivity_polars(
        df,
        insert=True,
        delete=True,
        update=True,
    )
    assert df.equals(original)


def test_empty_dataframe():
    df = pl.DataFrame({"id": pl.Series([], dtype=pl.Int64)})
    assert (
        alifestd_check_topological_sensitivity_polars(
            df,
            insert=True,
            delete=True,
            update=True,
        )
        == []
    )


def test_empty_dataframe_with_sensitive():
    df = pl.DataFrame(
        {
            "id": pl.Series([], dtype=pl.Int64),
            "branch_length": pl.Series([], dtype=pl.Float64),
        }
    )
    assert alifestd_check_topological_sensitivity_polars(
        df,
        insert=True,
        delete=True,
        update=True,
    ) == ["branch_length"]


@pytest.mark.parametrize("col", sorted(_topologically_sensitive_cols))
def test_lazyframe(base_df: pl.DataFrame, col: str):
    df = base_df.with_columns(pl.lit(0).alias(col))
    lazy = df.lazy()
    result = alifestd_check_topological_sensitivity_polars(
        lazy,
        insert=True,
        delete=True,
        update=True,
    )
    assert result == [col]


def test_lazyframe_none_present(base_df: pl.DataFrame):
    lazy = base_df.lazy()
    result = alifestd_check_topological_sensitivity_polars(
        lazy,
        insert=True,
        delete=True,
        update=True,
    )
    assert result == []


@pytest.mark.parametrize("col", sorted(_update_only_sensitive_cols))
def test_insert_only_excludes_update_only(base_df: pl.DataFrame, col: str):
    df = base_df.with_columns(pl.lit(0).alias(col))
    result = alifestd_check_topological_sensitivity_polars(
        df,
        insert=True,
        delete=False,
        update=False,
    )
    assert col not in result


@pytest.mark.parametrize(
    "col",
    sorted(_topologically_sensitive_cols - _update_only_sensitive_cols),
)
def test_insert_only_includes_structure_sensitive(
    base_df: pl.DataFrame, col: str
):
    df = base_df.with_columns(pl.lit(0).alias(col))
    result = alifestd_check_topological_sensitivity_polars(
        df,
        insert=True,
        delete=False,
        update=False,
    )
    assert result == [col]


@pytest.mark.parametrize("col", sorted(_update_only_sensitive_cols))
def test_delete_only_excludes_update_only(base_df: pl.DataFrame, col: str):
    df = base_df.with_columns(pl.lit(0).alias(col))
    result = alifestd_check_topological_sensitivity_polars(
        df,
        insert=False,
        delete=True,
        update=False,
    )
    assert col not in result


@pytest.mark.parametrize(
    "col",
    sorted(_topologically_sensitive_cols - _update_only_sensitive_cols),
)
def test_delete_only_includes_structure_sensitive(
    base_df: pl.DataFrame, col: str
):
    df = base_df.with_columns(pl.lit(0).alias(col))
    result = alifestd_check_topological_sensitivity_polars(
        df,
        insert=False,
        delete=True,
        update=False,
    )
    assert result == [col]


@pytest.mark.parametrize("col", sorted(_topologically_sensitive_cols))
def test_update_includes_all(base_df: pl.DataFrame, col: str):
    df = base_df.with_columns(pl.lit(0).alias(col))
    result = alifestd_check_topological_sensitivity_polars(
        df,
        insert=False,
        delete=False,
        update=True,
    )
    assert result == [col]


def test_no_ops_returns_empty(base_df: pl.DataFrame):
    cols = [pl.lit(0).alias(col) for col in _topologically_sensitive_cols]
    df = base_df.with_columns(cols)
    result = alifestd_check_topological_sensitivity_polars(
        df,
        insert=False,
        delete=False,
        update=False,
    )
    assert result == []


@pytest.mark.parametrize("col", sorted(_update_only_sensitive_cols))
def test_insert_only_lazyframe_excludes(base_df: pl.DataFrame, col: str):
    lazy = base_df.with_columns(pl.lit(0).alias(col)).lazy()
    result = alifestd_check_topological_sensitivity_polars(
        lazy,
        insert=True,
        delete=False,
        update=False,
    )
    assert col not in result
