import hashlib
import json
import os
import pathlib

import dask
import pandas as pd
import pytest
import toolz
import yaml

import xorq.api as xo
import xorq.vendor.ibis as ibis
import xorq.vendor.ibis.expr.operations.relations as rel
from xorq.caching import (
    ParquetCache,
    ParquetSnapshotCache,
    ParquetTTLSnapshotCache,
    SourceCache,
    SourceSnapshotCache,
)
from xorq.common.utils.dask_normalize.dask_normalize_utils import (
    normalize_read_path_md5sum,
)
from xorq.common.utils.defer_utils import deferred_read_parquet
from xorq.common.utils.graph_utils import find_all_sources, walk_nodes
from xorq.common.utils.name_utils import get_uid_prefix
from xorq.ibis_yaml.compiler import (
    ArtifactStore,
    DumpFiles,
    RefEnum,
    build_expr,
    load_expr,
)
from xorq.ibis_yaml.config import config
from xorq.ibis_yaml.sql import find_relations
from xorq.tests.util import assert_frame_equal
from xorq.vendor.ibis.common.collections import FrozenOrderedDict


do_roundtrip_expr = toolz.compose(load_expr, build_expr)


def get_local_path(parquet_dir, name):
    pins_path = parquet_dir.joinpath(f"{name}.parquet")
    local_path = pathlib.Path(pins_path.name)
    local_path.write_bytes(pins_path.read_bytes())
    return local_path


@pytest.mark.snapshot_check
def test_artifact_store_expr_hash(t, builds_dir, snapshot):
    artifact_store = ArtifactStore.from_path_and_expr(builds_dir, t)
    actual = artifact_store.root_path.name
    snapshot.assert_match(actual, "artifact_store_expr_hash.txt")


@pytest.fixture(scope="session")
def users_df():
    return pd.DataFrame(
        {
            "user_id": [1, 2, 3, 4, 5],
            "age": [25, 32, 28, 45, 31],
            "name": ["Alice", "Bob", "Charlie", "Diana", "Eve"],
        }
    )


def test_build_manager_roundtrip(t, builds_dir):
    artifact_store = ArtifactStore.from_path_and_expr(builds_dir, t)
    yaml_dict = {"a": "string"}
    artifact_store.save_yaml(yaml_dict, DumpFiles.expr)

    out = artifact_store.root_path.joinpath(DumpFiles.expr).read_text()
    assert out == "a: string\n"
    result = artifact_store.load_yaml(DumpFiles.expr)
    assert result == yaml_dict


def test_build_manager_paths(t, builds_dir):
    new_path = builds_dir / "new_path"

    assert not new_path.exists()
    ArtifactStore.from_path_and_expr(new_path, t)
    assert new_path.exists()


def test_clean_frozen_dict_yaml(builds_dir):
    artifact_store = ArtifactStore(builds_dir)
    data = FrozenOrderedDict(
        {"string": "text", "integer": 42, "float": 3.14, "boolean": True, "none": None}
    )

    expected_yaml = """string: text
integer: 42
float: 3.14
boolean: true
none: null
"""
    out_path = artifact_store.save_yaml(data, DumpFiles.expr)
    result = out_path.read_text()

    assert expected_yaml == result


def test_ibis_compiler(t, builds_dir):
    t = xo.memtable({"a": [0, 1], "b": [0, 1]})
    expr = t.filter(t.a == 1).drop("b")
    roundtrip_expr = do_roundtrip_expr(expr, builds_dir=builds_dir)
    assert expr.execute().equals(roundtrip_expr.execute())


def test_ibis_compiler_parquet_reader(builds_dir, parquet_dir):
    backend = xo.duckdb.connect()
    parquet_path = parquet_dir / "awards_players.parquet"
    awards_players = deferred_read_parquet(
        parquet_path, backend, table_name="award_players"
    )
    expr = awards_players.filter(awards_players.lgID == "NL").drop("yearID", "lgID")
    roundtrip_expr = do_roundtrip_expr(expr, builds_dir=builds_dir)
    assert expr.execute().equals(roundtrip_expr.execute())


def test_compiler_sql(builds_dir, parquet_dir):
    backend = xo.datafusion.connect()
    awards_players = deferred_read_parquet(
        parquet_dir / "awards_players.parquet",
        backend,
        table_name="awards_players",
    )
    expr = awards_players.filter(awards_players.lgID == "NL").drop("yearID", "lgID")

    build_path = build_expr(expr, builds_dir=builds_dir, debug=True)
    # make sure we can load
    load_expr(build_path)
    expected_relation = find_relations(awards_players)[0]
    expted_sql_hash = dask.base.tokenize(str(ibis.to_sql(expr)))[: config.hash_length]

    assert build_path.joinpath(DumpFiles.sql).exists()
    assert build_path.joinpath(DumpFiles.metadata).exists()
    metadata = json.loads(build_path.joinpath(DumpFiles.metadata).read_text())

    assert "current_library_version" in metadata
    sql_text = build_path.joinpath(DumpFiles.sql).read_text()
    expected_result = (
        "queries:\n"
        "  main:\n"
        "    engine: datafusion\n"
        f"    profile_name: {expr._find_backend()._profile.hash_name}\n"
        "    relations:\n"
        f"    - {expected_relation}\n"
        "    options: {}\n"
        f"    sql_file: {expted_sql_hash}.sql\n"
    )
    assert sql_text == expected_result


def test_deferred_reads_yaml(builds_dir, parquet_dir):
    backend = xo.datafusion.connect()
    # Factor out the config path
    config_path = parquet_dir / "awards_players.parquet"
    awards_players = deferred_read_parquet(
        config_path,
        backend,
        table_name="awards_players",
    )
    expr = awards_players.filter(awards_players.lgID == "NL").drop("yearID", "lgID")

    expected_relation = find_relations(awards_players)[0]
    expected_profile = backend._profile.hash_name

    build_path = build_expr(expr, builds_dir=builds_dir, debug=True)
    yaml_path = build_path.joinpath(DumpFiles.deferred_reads)
    assert yaml_path.exists()
    sql_text = yaml_path.read_text()

    sql_str = str(ibis.to_sql(awards_players))
    expected_sql_file = dask.base.tokenize(sql_str)[: config.hash_length] + ".sql"

    expected_read_path = str(config_path)

    expected_result = (
        "reads:\n"
        f"  {expected_relation}:\n"
        "    engine: datafusion\n"
        f"    profile_name: {expected_profile}\n"
        "    relations:\n"
        f"    - {expected_relation}\n"
        "    options:\n"
        "      method_name: read_parquet\n"
        "      name: awards_players\n"
        "      read_kwargs:\n"
        f"      - path: {expected_read_path}\n"
        "      - table_name: awards_players\n"
        f"    sql_file: {expected_sql_file}\n"
    )

    assert sql_text == expected_result


def test_ibis_compiler_expr_schema_ref(t, builds_dir):
    t = xo.memtable({"a": [0, 1], "b": [0, 1]})
    expr = t.filter(t.a == 1).drop("b")
    build_path = build_expr(expr, builds_dir=builds_dir)
    yaml_dict = yaml.safe_load(build_path.joinpath(DumpFiles.expr).read_text())
    assert yaml_dict["expression"][RefEnum.schema_ref]


def test_multi_engine_deferred_reads(builds_dir, parquet_dir):
    con0 = xo.connect()
    con1 = xo.connect()
    con2 = xo.duckdb.connect()
    con3 = xo.connect()

    awards_players = deferred_read_parquet(
        parquet_dir / "awards_players.parquet", con=con0
    ).into_backend(con1)
    batting = deferred_read_parquet(
        parquet_dir / "batting.parquet", con=con2
    ).into_backend(con1)
    expr = (
        awards_players.join(batting, predicates=["playerID", "yearID", "lgID"])
        .into_backend(con3)
        .filter(xo._.G == 1)
    )
    roundtrip_expr = do_roundtrip_expr(expr, builds_dir=builds_dir)
    assert expr.execute().equals(roundtrip_expr.execute())


def test_multi_engine_with_caching(builds_dir, parquet_dir):
    con0 = xo.connect()
    con1 = xo.connect()
    con2 = xo.duckdb.connect()
    con3 = xo.connect()

    awards_players = deferred_read_parquet(
        parquet_dir / "awards_players.parquet", con=con0
    ).into_backend(con1)
    batting = deferred_read_parquet(
        parquet_dir / "batting.parquet", con=con2
    ).into_backend(con1)
    expr = (
        awards_players.join(batting, predicates=["playerID", "yearID", "lgID"])
        .into_backend(con3)
        .filter(xo._.G == 1)
    )
    roundtrip_expr = do_roundtrip_expr(expr, builds_dir=builds_dir)
    assert expr.execute().equals(roundtrip_expr.execute())


@pytest.mark.parametrize(
    "environment_factory",
    (
        pytest.param(None, id="no_env"),
        pytest.param(lambda p: p / "env_cache", id="with_env"),
    ),
)
@pytest.mark.parametrize(
    "cli_factory",
    (
        pytest.param(None, id="no_cli"),
        pytest.param(lambda p: p / "cli_cache", id="with_cli"),
    ),
)
def test_multi_engine_with_caching_with_parquet(
    builds_dir, tmp_path, environment_factory, cli_factory, monkeypatch, parquet_dir
):
    con0 = xo.connect()
    con1 = xo.connect()

    cache = ParquetCache.from_kwargs(source=con1, relative_path=tmp_path)

    expr = (
        deferred_read_parquet(parquet_dir / "awards_players.parquet", con=con0)
        .into_backend(con1)
        .filter(xo._.playerID == "bondto01")
        .cache(cache=cache)
    )

    expected_cache_dir = tmp_path
    if environment_factory is not None:
        cache_dir = environment_factory(tmp_path)
        monkeypatch.setenv("XORQ_CACHE_DIR", str(cache_dir))
        expected_cache_dir = cache_dir.joinpath(tmp_path)

    if cli_factory is not None:
        cli_cache_dir = cli_factory(tmp_path)
        cache_dir = cli_cache_dir
        expected_cache_dir = cli_cache_dir.joinpath(tmp_path)
    else:
        cache_dir = None

    roundtrip_expr = load_expr(
        build_expr(expr, builds_dir=builds_dir, cache_dir=cache_dir)
    )
    assert expr.execute().equals(roundtrip_expr.execute())
    assert expected_cache_dir.exists()


@pytest.mark.parametrize(
    "table_from_df",
    (
        pytest.param(lambda _, df: xo.memtable(df, name="users"), id="memtable"),
        pytest.param(
            lambda con, df: con.register(df, table_name="users"), id="database_table"
        ),
    ),
)
def test_roundtrip_database_table(builds_dir, users_df, table_from_df):
    original = xo.connect()

    t = table_from_df(original, users_df)
    expr = t.filter(t.age > 30).select(t.user_id, t.name)
    roundtrip_expr = do_roundtrip_expr(expr, builds_dir=builds_dir)
    assert_frame_equal(xo.execute(expr), roundtrip_expr.execute())


@pytest.mark.parametrize(
    "table_from_df",
    (
        pytest.param(lambda _, df: xo.memtable(df, name="users"), id="memtable"),
        pytest.param(
            lambda con, df: con.register(df, table_name="users"), id="database_table"
        ),
    ),
)
def test_roundtrip_database_table_cached(builds_dir, tmp_path, users_df, table_from_df):
    original = xo.connect()
    ddb = xo.duckdb.connect()

    cache = ParquetCache.from_kwargs(source=ddb, relative_path=tmp_path)

    t = table_from_df(original, users_df)
    expr = t.filter(t.age > 30).select(t.user_id, t.name, t.age * 2).cache(cache=cache)
    roundtrip_expr = do_roundtrip_expr(expr, builds_dir=builds_dir)
    assert_frame_equal(xo.execute(expr), roundtrip_expr.execute())


@pytest.mark.parametrize(
    "table_from_df",
    (
        pytest.param(lambda _, df: xo.memtable(df, name="users"), id="memtable"),
        pytest.param(
            lambda con, df: con.register(df, table_name="users"), id="database_table"
        ),
    ),
)
def test_roundtrip_database_table_behind_cache(
    builds_dir, tmp_path, users_df, table_from_df
):
    original = xo.connect()
    ddb = xo.duckdb.connect()

    cache = ParquetCache.from_kwargs(source=ddb, relative_path=tmp_path)

    t = table_from_df(original, users_df)
    expr = (
        t.filter(t.age > 30)
        .cache(cache=cache)
        .select(xo._.user_id, xo._.name, xo._.age * 2)
    )
    roundtrip_expr = do_roundtrip_expr(expr, builds_dir=builds_dir)
    assert_frame_equal(xo.execute(expr), roundtrip_expr.execute())


def test_build_pandas_backend(builds_dir, users_df):
    xo_con = xo.connect()
    pandas_con = xo.pandas.connect()
    t = xo_con.register(users_df, table_name="users")

    expected = (
        t.filter(t.age > 30)
        .select(t.user_id, t.name, t.age * 2)
        .into_backend(pandas_con, name="pandas_users")
    )
    actual = do_roundtrip_expr(expected, builds_dir=builds_dir)
    assert_frame_equal(xo.execute(expected), actual.execute())


@pytest.mark.slow(level=1)
@pytest.mark.snapshot_check
def test_build_file_stability_https(builds_dir, snapshot):
    def with_profile_idx(con, idx):
        profile = con._profile
        con._profile = profile.clone(idx=idx)
        return con

    con0 = with_profile_idx(xo.connect(), 0)
    con1 = with_profile_idx(xo.connect(), 1)
    con2 = with_profile_idx(xo.duckdb.connect(), 2)
    con3 = with_profile_idx(xo.connect(), 3)

    awards_players_path = "https://storage.googleapis.com/letsql-pins/awards_players/20240711T171119Z-886c4/awards_players.parquet"
    batting_path = "https://storage.googleapis.com/letsql-pins/batting/20240711T171118Z-431ef/batting.parquet"

    awards_players = xo.deferred_read_parquet(
        awards_players_path,
        con0,
        "awards_players",
    ).into_backend(con1, "awards_players_into")
    batting = xo.deferred_read_parquet(
        batting_path,
        con2,
        "batting",
    ).into_backend(con1, "batting_into")
    expr = (
        awards_players.join(batting, predicates=["playerID", "yearID", "lgID"])
        .into_backend(con3, "joined_into")
        .filter(xo._.G == 1)
    )
    build_path = build_expr(expr, builds_dir=builds_dir, debug=True)

    actual = json.dumps(
        {
            p.name: hashlib.md5(p.read_bytes()).hexdigest()
            for p in build_path.iterdir()
            if p.name != DumpFiles.metadata
        },
        indent=2,
        sort_keys=True,
    )

    snapshot.assert_match(actual, "expected.json")

    # test that it also runs
    roundtrip_expr = load_expr(build_path)
    assert expr.execute().equals(roundtrip_expr.execute())


@pytest.mark.snapshot_check
def test_build_file_stability_local(
    builds_dir,
    parquet_dir,
    tmpdir,
    monkeypatch,
    snapshot,
):
    monkeypatch.chdir(tmpdir)

    def with_profile_idx(con, idx):
        profile = con._profile
        con._profile = profile.clone(idx=idx)
        return con

    batting_path = get_local_path(parquet_dir, "batting")
    awards_players_path = get_local_path(parquet_dir, "awards_players")

    con0 = with_profile_idx(xo.connect(), 0)
    con1 = with_profile_idx(xo.connect(), 1)
    con2 = with_profile_idx(xo.duckdb.connect(), 2)
    con3 = with_profile_idx(xo.connect(), 3)

    awards_players = xo.deferred_read_parquet(
        awards_players_path,
        con0,
        "awards_players",
        # we must hash based on content: inode stat is constantly updating
        normalize_method=normalize_read_path_md5sum,
    ).into_backend(con1, "awards_players_into")
    batting = xo.deferred_read_parquet(
        batting_path,
        con2,
        "batting",
        # we must hash based on content: inode stat is constantly updating
        normalize_method=normalize_read_path_md5sum,
    ).into_backend(con1, "batting_into")
    expr = (
        awards_players.join(batting, predicates=["playerID", "yearID", "lgID"])
        .into_backend(con3, "joined_into")
        .filter(xo._.G == 1)
    )
    build_path = build_expr(expr, builds_dir=builds_dir, debug=True)
    actual = json.dumps(
        {
            p.name: hashlib.md5(p.read_bytes()).hexdigest()
            for p in build_path.iterdir()
            if p.name != DumpFiles.metadata
        },
        indent=2,
        sort_keys=True,
    )

    snapshot.assert_match(actual, "expected.json")

    # test that it also runs
    roundtrip_expr = load_expr(build_path)
    assert expr.execute().equals(roundtrip_expr.execute())


@pytest.mark.snapshot_check
def test_build_file_stability_and_relocatability(
    builds_dir, parquet_dir, tmpdir, monkeypatch, snapshot
):
    # path impacts node hash therefore path ***MUST*** be the same
    monkeypatch.chdir(tmpdir)

    path = get_local_path(parquet_dir, "awards_players")
    awards_players = xo.deferred_read_parquet(
        path,
        normalize_method=normalize_read_path_md5sum,
    )
    batting = xo.memtable(
        xo.read_parquet(parquet_dir.joinpath("batting.parquet")).execute()
    )
    on = sorted(set(batting.columns).intersection(awards_players.columns))
    expr = awards_players.select(on).join(batting.select(on), predicates=on)

    build_dir = build_expr(
        expr, builds_dir=builds_dir, read_normalize_method=normalize_read_path_md5sum
    )
    actual = json.dumps(
        {
            p.name: hashlib.md5(p.read_bytes()).hexdigest()
            for p in build_dir.iterdir()
            if p.name != DumpFiles.metadata and p.is_file()
        }
        | {
            "build_dir_name": build_dir.name,
        },
        indent=2,
        sort_keys=True,
    )
    snapshot.assert_match(actual, "expected.json")

    # test that it also runs
    roundtrip_expr = load_expr(build_dir)
    (actual, expected) = (
        expr.execute().pipe(lambda t: t.sort_values(list(t.columns), ignore_index=True))
        for expr in (roundtrip_expr, expr)
    )
    assert actual.equals(expected)

    before = tuple(node.name for node in walk_nodes(rel.PhysicalTable, expr))
    assert all(map(get_uid_prefix, before))
    after = tuple(node.name for node in walk_nodes(rel.PhysicalTable, roundtrip_expr))
    assert not any(map(get_uid_prefix, after))


def test_build_pandas_backend_behind_into_backend(builds_dir, users_df):
    xo_con = xo.connect()
    pandas_con = xo.pandas.connect()
    t = xo_con.register(users_df, table_name="users")

    expected = (
        t.filter(t.age > 30)
        .into_backend(pandas_con, name="pandas_users")
        .select(xo._.user_id, xo._.name, xo._.age * 2)
    )
    actual = do_roundtrip_expr(expected, builds_dir=builds_dir)
    assert_frame_equal(xo.execute(expected), actual.execute())


def test_struct_field(builds_dir, tmpdir):
    path = pathlib.Path(tmpdir).joinpath("t.parquet")
    xo.memtable({"a": [{"b": 1, "c": "string"}]}).to_parquet(path)
    t = xo.deferred_read_parquet(
        path,
        xo.connect(),
        table_name="t",
    )
    expr = t.select(t.a.b.name("a-b"))
    roundtrip_expr = do_roundtrip_expr(expr, builds_dir=builds_dir)
    assert_frame_equal(expr.execute(), roundtrip_expr.execute())


def test_no_sql_or_deferred_when_debug_false(builds_dir):
    t = xo.memtable({"a": [1, 2, 3]})
    expr = t.filter(t.a > 1)
    build_path = build_expr(expr, builds_dir=builds_dir, debug=False)
    assert not os.path.exists(build_path / DumpFiles.sql)
    assert not os.path.exists(build_path / DumpFiles.deferred_reads)


def test_into_backend_with_array_filter(builds_dir):
    from xorq.conftest import array_types_df

    duckdb_con = xo.duckdb.connect()

    t = duckdb_con.create_table("array_types", array_types_df)
    expr = t.mutate(filtered=t.x.filter(xo._ > 1)).cache(
        SourceCache.from_kwargs(source=xo.connect())
    )
    roundtrip_expr = do_roundtrip_expr(expr, builds_dir=builds_dir, debug=False)
    assert_frame_equal(expr.execute(), roundtrip_expr.execute())
    assert {"duckdb", "xorq"}.intersection(
        source.name for source in find_all_sources(roundtrip_expr)
    )


def test_roundtrip_parquet_snapshot_cache(builds_dir, tmp_path, users_df):
    original = xo.connect()
    ddb = xo.duckdb.connect()

    cache = ParquetSnapshotCache.from_kwargs(source=ddb, relative_path=tmp_path)

    t = original.register(users_df, table_name="users")
    expr = t.filter(t.age > 30).select(t.user_id, t.name, t.age * 2).cache(cache=cache)
    roundtrip_expr = do_roundtrip_expr(expr, builds_dir=builds_dir)
    assert_frame_equal(xo.execute(expr), roundtrip_expr.execute())


def test_roundtrip_parquet_ttl_snapshot_cache(builds_dir, tmp_path, users_df):
    import datetime

    original = xo.connect()
    ddb = xo.duckdb.connect()

    cache = ParquetTTLSnapshotCache.from_kwargs(
        source=ddb, relative_path=tmp_path, ttl=datetime.timedelta(hours=2)
    )

    t = original.register(users_df, table_name="users")
    expr = t.filter(t.age > 30).select(t.user_id, t.name, t.age * 2).cache(cache=cache)
    roundtrip_expr = do_roundtrip_expr(expr, builds_dir=builds_dir)
    assert_frame_equal(xo.execute(expr), roundtrip_expr.execute())


def test_roundtrip_source_snapshot_cache(builds_dir, users_df):
    original = xo.connect()

    cache = SourceSnapshotCache.from_kwargs(source=xo.connect())

    t = original.register(users_df, table_name="users")
    expr = t.filter(t.age > 30).select(t.user_id, t.name, t.age * 2).cache(cache=cache)
    roundtrip_expr = do_roundtrip_expr(expr, builds_dir=builds_dir)
    assert_frame_equal(xo.execute(expr), roundtrip_expr.execute())


def test_generated_name_sanitization_parquet(
    builds_dir,
    parquet_dir,
    tmpdir,
    monkeypatch,
    snapshot,
):
    monkeypatch.chdir(tmpdir)

    batting_path = get_local_path(parquet_dir, "batting")
    expr = xo.deferred_read_parquet(batting_path)
    build_path = build_expr(
        expr, builds_dir=builds_dir, read_normalize_method=normalize_read_path_md5sum
    )
    loaded = load_expr(build_path)

    assert (expr_name := expr.op().name) != (build_name := loaded.op().name)
    assert get_uid_prefix(expr_name)
    assert not get_uid_prefix(build_name)
    snapshot.assert_match(build_name, "parquet-build-name.txt")


def test_generated_name_sanitization_memtable(
    builds_dir,
    parquet_dir,
    snapshot,
):
    df = xo.deferred_read_parquet(parquet_dir.joinpath("batting.parquet")).execute()
    expr = xo.memtable(df.sort_values(list(df.columns)))
    build_path = build_expr(expr, builds_dir=builds_dir)
    loaded = load_expr(build_path)

    assert (expr_name := expr.op().name) != (build_name := loaded.op().name)
    assert get_uid_prefix(expr_name)
    assert not get_uid_prefix(build_name)
    snapshot.assert_match(build_name, "memory-build-name.txt")


def test_pandas_memtable_comparison(builds_dir):
    df = pd.DataFrame({"a": [1]})
    expr = xo.memtable(df, name="name")
    expr2 = xo.memtable(df, name="name")
    joined = expr.join(expr2, predicates="a")
    xo.build_expr(joined, builds_dir=builds_dir)


def test_pyarrow_memtable_comparison(builds_dir):
    import pyarrow as pa

    table = pa.table({"a": [1]})
    expr = xo.memtable(table, name="name")
    expr2 = xo.memtable(table, name="name")
    joined = expr.join(expr2, predicates="a")
    xo.build_expr(joined, builds_dir=builds_dir)


def test_polars_memtable_comparison(builds_dir):
    pl = pytest.importorskip("polars")

    df = pl.DataFrame({"a": [1]})
    expr = xo.memtable(df, name="name")
    expr2 = xo.memtable(df, name="name")
    joined = expr.join(expr2, predicates="a")
    xo.build_expr(joined, builds_dir=builds_dir)
