import pandas as pd
import pyarrow.compute as pc
import pytest

import xorq.api as xo
import xorq.expr.datatypes as dt
import xorq.expr.relations as rel
from xorq.caching import ParquetCache, SourceCache
from xorq.common.utils.graph_utils import find_all_sources, replace_sources


def assert_result_equal(left, right):
    """Compare DataFrames ignoring row order (backends may return different orderings)."""
    left = left.sort_values(left.columns.tolist()).reset_index(drop=True)
    right = right.sort_values(right.columns.tolist()).reset_index(drop=True)
    pd.testing.assert_frame_equal(left, right)


@xo.udf.scalar.pyarrow
def double_int(arr: dt.int64) -> dt.int64:
    return pc.multiply(arr, 2)


@pytest.fixture(scope="session")
def parquet_path(parquet_dir):
    return parquet_dir / "batting.parquet"


BACKEND_NAMES = ("xorq", "duckdb", "pandas")


def make_con(name):
    return getattr(xo, name).connect()


def backend_pair_id(val):
    return f"{val[0]}->{val[1]}"


BACKEND_PAIRS = tuple(
    (from_name, to_name) for from_name in BACKEND_NAMES for to_name in BACKEND_NAMES
)


@pytest.fixture(params=BACKEND_PAIRS, ids=[backend_pair_id(p) for p in BACKEND_PAIRS])
def from_to(request):
    return request.param


# ---------------------------------------------------------------------------
# Parametrized: registered table (requires transfer_tables=True)
# ---------------------------------------------------------------------------


def test_replace_registered_table(parquet_path, from_to):
    """Replace source on a table registered via read_parquet."""
    from_name, to_name = from_to
    from_con = make_con(from_name)
    to_con = make_con(to_name)

    t = from_con.read_parquet(parquet_path, table_name="batting")
    result = replace_sources({id(from_con): to_con}, t, transfer_tables=True)

    sources = find_all_sources(result)
    assert len(sources) == 1
    assert sources[0] is to_con
    assert_result_equal(result.execute(), t.execute())


def test_replace_registered_table_raises_without_transfer(parquet_path):
    """Replacing a registered table without transfer_tables raises."""
    xo_con = xo.connect()
    t = xo_con.read_parquet(parquet_path, table_name="batting")

    ddb_con = xo.duckdb.connect()
    with pytest.raises(ValueError, match="DatabaseTable"):
        replace_sources({id(xo_con): ddb_con}, t)


# ---------------------------------------------------------------------------
# Parametrized: deferred read (no transfer needed)
# ---------------------------------------------------------------------------


def test_replace_deferred_read(parquet_path, from_to):
    """Replace source on a deferred_read_parquet expression."""
    from_name, to_name = from_to
    from_con = make_con(from_name)
    to_con = make_con(to_name)

    t = xo.deferred_read_parquet(parquet_path, from_con, table_name="dr_batting")
    result = replace_sources({id(from_con): to_con}, t)

    sources = find_all_sources(result)
    assert len(sources) == 1
    assert sources[0] is to_con
    assert isinstance(result.op(), rel.Read)
    assert result.op().source is to_con
    assert_result_equal(result.execute(), t.execute())


# ---------------------------------------------------------------------------
# Parametrized: projection and filter (requires transfer_tables=True)
# ---------------------------------------------------------------------------


def test_replace_with_projection(parquet_path, from_to):
    """Replace source on an expression with column selections and filters."""
    from_name, to_name = from_to
    from_con = make_con(from_name)
    to_con = make_con(to_name)

    t = from_con.read_parquet(parquet_path, table_name="batting")
    expr = t.filter(t.yearID > 2000).select("playerID", "yearID", "G")

    result = replace_sources({id(from_con): to_con}, expr, transfer_tables=True)

    sources = find_all_sources(result)
    assert len(sources) == 1
    assert sources[0] is to_con
    assert_result_equal(result.execute(), expr.execute())


# ---------------------------------------------------------------------------
# Parametrized: into_backend
# ---------------------------------------------------------------------------


def test_replace_into_backend_target(parquet_path, from_to):
    """Replace the target backend of an into_backend expression."""
    from_name, to_name = from_to
    if from_name == "pandas" or to_name == "pandas":
        pytest.skip("pandas does not support into_backend")

    source_con = make_con(from_name)
    target_con = make_con(to_name)

    t = source_con.read_parquet(parquet_path, table_name="batting")
    expr = t.into_backend(target_con)

    replacement_con = make_con(to_name)
    result = replace_sources({id(target_con): replacement_con}, expr)

    sources = find_all_sources(result)
    assert replacement_con in sources
    assert target_con not in sources
    assert_result_equal(result.execute(), expr.execute())


# ---------------------------------------------------------------------------
# Parametrized: aggregation (requires transfer_tables=True)
# ---------------------------------------------------------------------------


def test_replace_with_aggregation(parquet_path, from_to):
    """Replace source on an expression with aggregation."""
    from_name, to_name = from_to
    from_con = make_con(from_name)
    to_con = make_con(to_name)

    t = from_con.read_parquet(parquet_path, table_name="batting")
    expr = t.group_by("yearID").agg(total_G=t.G.sum())

    result = replace_sources({id(from_con): to_con}, expr, transfer_tables=True)

    sources = find_all_sources(result)
    assert len(sources) == 1
    assert sources[0] is to_con
    assert_result_equal(result.execute(), expr.execute())


# ---------------------------------------------------------------------------
# Parametrized: self-join (requires transfer_tables=True)
# ---------------------------------------------------------------------------


def test_replace_join_same_backend(parquet_path, from_to):
    """Replace source in a self-join on the same backend."""
    from_name, to_name = from_to
    from_con = make_con(from_name)
    to_con = make_con(to_name)

    t = from_con.read_parquet(parquet_path, table_name="batting")
    expr = t.join(t, predicates=["playerID", "yearID"])

    result = replace_sources({id(from_con): to_con}, expr, transfer_tables=True)

    sources = find_all_sources(result)
    assert all(s is to_con for s in sources)
    assert_result_equal(result.execute(), expr.execute())


# ---------------------------------------------------------------------------
# Parametrized: scalar UDF (requires transfer_tables=True)
# ---------------------------------------------------------------------------


def test_replace_with_scalar_udf(parquet_path, from_to):
    """Replace source on an expression that applies a scalar UDF."""
    from_name, to_name = from_to
    if "pandas" in (from_name, to_name):
        pytest.skip("pandas cannot execute pyarrow scalar UDFs")
    from_con = make_con(from_name)
    to_con = make_con(to_name)

    t = from_con.read_parquet(parquet_path, table_name="batting")
    expr = t.mutate(doubled_G=double_int(t.G))

    result = replace_sources({id(from_con): to_con}, expr, transfer_tables=True)

    sources = find_all_sources(result)
    assert len(sources) == 1
    assert sources[0] is to_con
    assert_result_equal(result.execute(), expr.execute())


# ---------------------------------------------------------------------------
# Caching (xorq-specific, not parametrized; requires transfer_tables=True)
# ---------------------------------------------------------------------------


def test_replace_source_with_source_cache(parquet_path):
    """Replace source when expression uses SourceCache."""
    xo_con = xo.connect()
    t = xo_con.read_parquet(parquet_path, table_name="batting")
    expr = t.cache(SourceCache.from_kwargs(source=xo_con))

    ddb_con = xo.duckdb.connect()
    result = replace_sources({id(xo_con): ddb_con}, expr, transfer_tables=True)

    sources = find_all_sources(result)
    assert all(s is ddb_con for s in sources)
    cache_op = result.op()
    assert isinstance(cache_op, rel.CachedNode)
    assert cache_op.source is ddb_con
    assert_result_equal(result.execute(), expr.execute())


def test_replace_source_with_parquet_cache(parquet_path):
    """Replace source when expression uses ParquetCache."""
    xo_con = xo.connect()
    t = xo_con.read_parquet(parquet_path, table_name="batting")
    expr = t.filter(t.yearID > 2000).cache(ParquetCache.from_kwargs(source=xo_con))

    ddb_con = xo.duckdb.connect()
    result = replace_sources({id(xo_con): ddb_con}, expr, transfer_tables=True)

    sources = find_all_sources(result)
    assert ddb_con in sources
    assert xo_con not in sources
    assert_result_equal(result.execute(), expr.execute())


def test_replace_source_with_chained_cache(parquet_path):
    """Replace source through a chain of cached expressions."""
    xo_con = xo.connect()
    t = xo_con.read_parquet(parquet_path, table_name="batting")
    step1 = t.filter(t.yearID > 2000).cache()
    expr = step1.filter(step1.G > 10).cache()

    ddb_con = xo.duckdb.connect()
    result = replace_sources({id(xo_con): ddb_con}, expr, transfer_tables=True)

    sources = find_all_sources(result)
    assert xo_con not in sources
    assert_result_equal(result.execute(), expr.execute())


# ---------------------------------------------------------------------------
# Multi-backend (not parametrized — fixed backend combinations)
# ---------------------------------------------------------------------------


def test_replace_one_of_two_backends(parquet_path):
    """Replace only one backend in a multi-source expression."""
    xo_con = xo.connect()
    ddb_con = xo.duckdb.connect()

    t1 = xo_con.read_parquet(parquet_path, table_name="batting")
    ddb_con.read_parquet(parquet_path, table_name="batting")
    t2 = ddb_con.table("batting")

    expr = t1.into_backend(ddb_con).join(t2, predicates=["playerID", "yearID"])

    ddb_con2 = xo.duckdb.connect()
    result = replace_sources({id(ddb_con): ddb_con2}, expr, transfer_tables=True)

    sources = find_all_sources(result)
    assert ddb_con not in sources
    assert xo_con in sources
    assert ddb_con2 in sources
    assert_result_equal(result.execute(), expr.execute())


def test_replace_both_backends(parquet_path):
    """Replace both backends in a multi-source expression."""
    xo_con = xo.connect()
    ddb_con = xo.duckdb.connect()

    t1 = xo_con.read_parquet(parquet_path, table_name="batting")
    ddb_con.read_parquet(parquet_path, table_name="batting")
    t2 = ddb_con.table("batting")

    expr = t1.into_backend(ddb_con).join(t2, predicates=["playerID", "yearID"])

    xo_con2 = xo.connect()
    ddb_con2 = xo.duckdb.connect()
    result = replace_sources(
        {id(xo_con): xo_con2, id(ddb_con): ddb_con2}, expr, transfer_tables=True
    )

    sources = find_all_sources(result)
    assert xo_con not in sources
    assert ddb_con not in sources
    assert xo_con2 in sources
    assert ddb_con2 in sources
    assert_result_equal(result.execute(), expr.execute())


def test_replace_source_in_into_backend_chain(parquet_path):
    """Replace source through xorq -> duckdb -> xorq chain."""
    xo_con1 = xo.connect()
    ddb_con = xo.duckdb.connect()
    xo_con2 = xo.connect()

    t = xo_con1.read_parquet(parquet_path, table_name="batting")
    step1 = t.into_backend(ddb_con)
    expr = step1.into_backend(xo_con2)

    xo_con3 = xo.connect()
    result = replace_sources({id(xo_con2): xo_con3}, expr)

    sources = find_all_sources(result)
    assert xo_con3 in sources
    assert xo_con2 not in sources
    assert xo_con1 in sources
    assert_result_equal(result.execute(), expr.execute())


# ---------------------------------------------------------------------------
# Edge cases (not parametrized)
# ---------------------------------------------------------------------------


def test_replace_sources_empty_mapping(parquet_path):
    """Empty mapping returns an equivalent expression."""
    xo_con = xo.connect()
    t = xo_con.read_parquet(parquet_path, table_name="batting")
    expr = t.filter(t.yearID > 2000)

    result = replace_sources({}, expr)
    assert result.equals(expr)
    assert_result_equal(result.execute(), expr.execute())


def test_replace_sources_no_matching_source(parquet_path):
    """Mapping that doesn't match any source returns equivalent expression."""
    xo_con = xo.connect()
    t = xo_con.read_parquet(parquet_path, table_name="batting")

    ddb_con = xo.duckdb.connect()
    result = replace_sources({id(ddb_con): xo.connect()}, t)
    assert result.equals(t)
    assert_result_equal(result.execute(), t.execute())


def test_replace_sources_preserves_schema(parquet_path):
    """Replacing source preserves the expression schema."""
    xo_con = xo.connect()
    t = xo_con.read_parquet(parquet_path, table_name="batting")
    original_schema = t.schema()

    ddb_con = xo.duckdb.connect()
    result = replace_sources({id(xo_con): ddb_con}, t, transfer_tables=True)

    assert result.schema() == original_schema
    assert_result_equal(result.execute(), t.execute())


def test_replace_sources_preserves_columns(parquet_path):
    """Replacing source preserves column names."""
    xo_con = xo.connect()
    t = xo_con.read_parquet(parquet_path, table_name="batting")
    original_columns = t.columns

    ddb_con = xo.duckdb.connect()
    result = replace_sources({id(xo_con): ddb_con}, t, transfer_tables=True)

    assert result.columns == original_columns
    assert_result_equal(result.execute(), t.execute())


def test_replace_sources_does_not_mutate_original(parquet_path):
    """Replacing source does not mutate the original expression."""
    xo_con = xo.connect()
    t = xo_con.read_parquet(parquet_path, table_name="batting")
    original_sources = find_all_sources(t)

    ddb_con = xo.duckdb.connect()
    replace_sources({id(xo_con): ddb_con}, t, transfer_tables=True)

    assert find_all_sources(t) == original_sources
    assert original_sources[0] is xo_con
