"""Reusable dataset compare/diff helpers shared by CLI and TUI commands."""

from __future__ import annotations

from collections.abc import Iterable, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any

import polars as pl
from polars import datatypes as pl_datatypes

from .core.engine.polars_adapter import collect_lazyframe, unwrap_physical_plan
from .data.db import is_db_uri

if TYPE_CHECKING:  # pragma: no cover - import for typing only
    from .api.runtime import Runtime


@dataclass(frozen=True, slots=True)
class CompareResult:
    """Result payload returned by :func:`compare_paths`."""

    diff_lf: pl.LazyFrame
    summary: dict[str, object]
    warnings: tuple[str, ...] = ()


def compare_paths(
    runtime: Runtime | Any,
    left_path: str | Path,
    right_path: str | Path,
    *,
    keys: Sequence[str] | None = None,
    cols: Sequence[str] | None = None,
    exclude: Sequence[str] | None = None,
    rtol: float | None = None,
    atol: float | None = None,
) -> CompareResult:
    """Compare two dataset paths using the runtime's scanners."""

    left = _validate_path(left_path)
    right = _validate_path(right_path)
    left_lf, left_schema = _load_lazyframe(runtime, left)
    right_lf, right_schema = _load_lazyframe(runtime, right)
    return compare_lazyframes(
        left_lf,
        right_lf,
        left_schema=left_schema,
        right_schema=right_schema,
        left_path=left,
        right_path=right,
        keys=list(keys or ()),
        cols=list(cols or ()),
        exclude=list(exclude or ()),
        rtol=rtol,
        atol=atol,
    )


def compare_lazyframes(
    left_lf: pl.LazyFrame,
    right_lf: pl.LazyFrame,
    *,
    left_schema: pl.Schema,
    right_schema: pl.Schema,
    left_path: Path,
    right_path: Path,
    keys: list[str],
    cols: list[str],
    exclude: list[str],
    rtol: float | None,
    atol: float | None,
) -> CompareResult:
    """Compare two lazyframes and return the diff lazyframe + summary."""

    if rtol is not None and rtol < 0:
        raise ValueError("--rtol must be non-negative")
    if atol is not None and atol < 0:
        raise ValueError("--atol must be non-negative")

    left_cols = list(left_schema.keys())
    right_cols = list(right_schema.keys())

    if keys and exclude:
        overlap = sorted(set(keys) & set(exclude))
        if overlap:
            raise ValueError(f"Excluded columns cannot include key columns: {', '.join(overlap)}")

    if keys:
        _assert_columns_present(keys, left_schema, "left")
        _assert_columns_present(keys, right_schema, "right")
        row_mode = False
        left_lf_final = left_lf
        right_lf_final = right_lf
    else:
        row_col = _unique_column_name("__pulka_row", left_cols, right_cols)
        keys = [row_col]
        row_mode = True
        left_lf_final = left_lf.with_row_index(name=row_col)
        right_lf_final = right_lf.with_row_index(name=row_col)

    left_extra, right_extra, compare_cols = _resolve_compare_columns(
        cols, exclude, left_cols, right_cols
    )
    if not compare_cols:
        raise ValueError("No comparable columns after exclusion.")

    dtype_mismatches = _dtype_mismatches(compare_cols, left_schema, right_schema)
    if dtype_mismatches:
        mismatch_text = "; ".join(dtype_mismatches)
        raise ValueError(f"dtype mismatch: {mismatch_text}")

    warnings: list[str] = []
    if not row_mode:
        left_nulls = _null_key_columns(left_lf_final, keys)
        right_nulls = _null_key_columns(right_lf_final, keys)
        if left_nulls:
            joined = ", ".join(left_nulls)
            warnings.append(
                f"left key columns contain nulls ({joined}); null keys may not match across files"
            )
        if right_nulls:
            joined = ", ".join(right_nulls)
            warnings.append(
                f"right key columns contain nulls ({joined}); null keys may not match across files"
            )
        _assert_unique_keys(left_lf_final, keys, "left")
        _assert_unique_keys(right_lf_final, keys, "right")

    diff_lf = _build_diff_lazyframe(
        left_lf_final,
        right_lf_final,
        keys=keys,
        compare_cols=compare_cols,
        dtype_map=left_schema,
        rtol=rtol,
        atol=atol,
    )

    left_rows = _row_count(left_lf_final)
    right_rows = _row_count(right_lf_final)
    summary = _summarize_diff(
        diff_lf,
        keys=keys,
        compare_cols=compare_cols,
        dtype_map=left_schema,
        left_extra=left_extra,
        right_extra=right_extra,
        rtol=rtol,
        atol=atol,
        left_path=left_path,
        right_path=right_path,
        row_mode=row_mode,
        left_rows=left_rows,
        right_rows=right_rows,
    )
    if warnings:
        summary = dict(summary)
        summary["warnings"] = tuple(warnings)
    if exclude:
        summary = dict(summary)
        summary["exclude"] = list(exclude)
    return CompareResult(diff_lf=diff_lf, summary=summary, warnings=tuple(warnings))


def _validate_path(value: str | Path) -> Path:
    text = str(value)
    if is_db_uri(text):
        raise ValueError("Database URIs are not supported by pulka compare.")
    path = Path(value).expanduser()
    if not path.exists():
        raise ValueError(f"Path not found: {value}")
    return path


def _load_lazyframe(runtime: Any, path: Path) -> tuple[pl.LazyFrame, pl.Schema]:
    physical_plan = runtime.scanners.scan(path)
    lazyframe = unwrap_physical_plan(physical_plan).to_lazyframe()
    try:
        schema = lazyframe.collect_schema()
    except Exception:
        schema = lazyframe.schema
    return lazyframe, schema


def parse_name_list(raw: str | None) -> list[str]:
    """Parse a comma-separated column name list, preserving order and uniqueness."""

    if not raw:
        return []
    values: list[str] = []
    for chunk in raw.split(","):
        item = chunk.strip()
        if not item:
            continue
        values.append(item)
    # Preserve first-seen order while removing duplicates.
    return list(dict.fromkeys(values))


def _assert_columns_present(columns: Iterable[str], schema: pl.Schema, label: str) -> None:
    missing = [name for name in columns if name not in schema]
    if missing:
        missing_text = ", ".join(missing)
        raise ValueError(f"Missing columns in {label} dataset: {missing_text}")


def _resolve_compare_columns(
    columns: list[str],
    exclude: list[str],
    left_cols: list[str],
    right_cols: list[str],
) -> tuple[list[str], list[str], list[str]]:
    left_set = set(left_cols)
    right_set = set(right_cols)
    left_extra = sorted(left_set - right_set)
    right_extra = sorted(right_set - left_set)

    if exclude:
        missing_exclude = [
            name for name in exclude if name not in left_set and name not in right_set
        ]
        if missing_exclude:
            missing_text = ", ".join(sorted(missing_exclude))
            raise ValueError(f"Unknown columns for exclusion: {missing_text}")
        exclude_set = set(exclude)
    else:
        exclude_set = set()

    if columns:
        left_missing = [name for name in columns if name not in left_set]
        right_missing = [name for name in columns if name not in right_set]
        missing = sorted(set(left_missing + right_missing))
        if missing:
            raise ValueError(f"Unknown columns for comparison: {', '.join(missing)}")
        compare_cols = [name for name in columns if name not in exclude_set]
        return left_extra, right_extra, compare_cols

    compare_cols = sorted(left_set & right_set)
    if exclude_set:
        compare_cols = [name for name in compare_cols if name not in exclude_set]
    return left_extra, right_extra, compare_cols


def _dtype_mismatches(
    columns: Iterable[str],
    left_schema: pl.Schema,
    right_schema: pl.Schema,
) -> list[str]:
    mismatches: list[str] = []
    for name in columns:
        left_dtype = left_schema.get(name)
        right_dtype = right_schema.get(name)
        if left_dtype != right_dtype:
            mismatches.append(f"{name} ({left_dtype} != {right_dtype})")
    return mismatches


def _null_key_columns(lf: pl.LazyFrame, keys: list[str]) -> list[str]:
    exprs = [pl.col(name).is_null().any().alias(name) for name in keys]
    result = collect_lazyframe(lf.select(exprs))
    if result.height == 0:
        return []
    row = result.row(0)
    return [name for name, flag in zip(keys, row, strict=True) if bool(flag)]


def _assert_unique_keys(lf: pl.LazyFrame, keys: list[str], label: str) -> None:
    dupes = collect_lazyframe(
        lf.select(keys).group_by(keys).len().filter(pl.col("len") > 1).limit(1)
    )
    if dupes.height:
        key_text = ", ".join(keys)
        raise ValueError(f"Duplicate keys detected in {label} dataset for [{key_text}].")


def _unique_column_name(base: str, *columns: Iterable[str]) -> str:
    taken = set()
    for col_group in columns:
        taken.update(col_group)
    if base not in taken:
        return base
    idx = 1
    while f"{base}_{idx}" in taken:
        idx += 1
    return f"{base}_{idx}"


def _is_float_dtype(dtype: pl.DataType) -> bool:
    is_float = getattr(pl_datatypes, "is_float", None)
    if callable(is_float):
        return bool(is_float(dtype))
    return dtype in (pl.Float32, pl.Float64)


def _float_diff_expr(
    left_col: pl.Expr,
    right_col: pl.Expr,
    *,
    rtol: float | None,
    atol: float | None,
) -> pl.Expr:
    """Return a diff expression for float columns with null==null semantics."""

    rtol_value = float(rtol or 0.0)
    atol_value = float(atol or 0.0)
    both_null = left_col.is_null() & right_col.is_null()
    both_value = left_col.is_not_null() & right_col.is_not_null()
    max_abs = pl.max_horizontal(left_col.abs(), right_col.abs())
    tol = pl.max_horizontal(pl.lit(atol_value), max_abs * rtol_value)
    approx_equal = (left_col - right_col).abs() <= tol
    equal_expr = both_null | (both_value & approx_equal)
    return equal_expr.not_().fill_null(True)


def _build_diff_lazyframe(
    left_lf: pl.LazyFrame,
    right_lf: pl.LazyFrame,
    *,
    keys: list[str],
    compare_cols: list[str],
    dtype_map: pl.Schema,
    rtol: float | None,
    atol: float | None,
) -> pl.LazyFrame:
    left_renames = [pl.col(name).alias(f"{name}__left") for name in compare_cols]
    right_renames = [pl.col(name).alias(f"{name}__right") for name in compare_cols]

    left_sel = [*map(pl.col, keys), *left_renames, pl.lit(True).alias("__present_left")]
    right_sel = [*map(pl.col, keys), *right_renames, pl.lit(True).alias("__present_right")]

    joined = left_lf.select(left_sel).join(
        right_lf.select(right_sel),
        on=keys,
        how="full",
        coalesce=True,
    )

    diff_exprs: list[pl.Expr] = []
    diff_exprs_by_col: dict[str, pl.Expr] = {}
    for name in compare_cols:
        left_name = f"{name}__left"
        right_name = f"{name}__right"
        left_col = pl.col(left_name)
        right_col = pl.col(right_name)
        if (rtol is not None or atol is not None) and _is_float_dtype(dtype_map.get(name)):
            diff_expr = _float_diff_expr(
                left_col,
                right_col,
                rtol=rtol,
                atol=atol,
            )
        else:
            diff_expr = left_col.eq_missing(right_col).not_()
        diff_exprs.append(diff_expr)
        diff_exprs_by_col[name] = diff_expr

    any_diff = pl.any_horizontal(*diff_exprs) if diff_exprs else pl.lit(False)
    present_left = pl.col("__present_left").is_not_null()
    present_right = pl.col("__present_right").is_not_null()

    global_kind_expr = (
        pl.when(~present_left & present_right)
        .then(pl.lit("right_only"))
        .when(~present_right & present_left)
        .then(pl.lit("left_only"))
        .when(any_diff)
        .then(pl.lit("changed"))
        .otherwise(pl.lit("unchanged"))
    )

    kind_exprs: list[pl.Expr] = []
    for name in compare_cols:
        kind_exprs.append(
            pl.when(~present_left & present_right)
            .then(pl.lit("right_only"))
            .when(~present_right & present_left)
            .then(pl.lit("left_only"))
            .when(diff_exprs_by_col[name])
            .then(pl.lit("changed"))
            .otherwise(pl.lit("unchanged"))
            .alias(f"{name}__kind")
        )

    ordered_cols = [*keys, "global__kind"]
    for name in compare_cols:
        ordered_cols.append(f"{name}__kind")
        ordered_cols.append(f"{name}__left")
        ordered_cols.append(f"{name}__right")

    return (
        joined.with_columns(global_kind_expr.alias("global__kind"), *kind_exprs)
        .select(ordered_cols)
        .sort(by=keys, nulls_last=True)
    )


def _summarize_diff(
    diff_lf: pl.LazyFrame,
    *,
    keys: list[str],
    compare_cols: list[str],
    dtype_map: pl.Schema,
    left_extra: list[str],
    right_extra: list[str],
    rtol: float | None,
    atol: float | None,
    left_path: Path,
    right_path: Path,
    row_mode: bool,
    left_rows: int,
    right_rows: int,
) -> dict[str, object]:
    counts_df = collect_lazyframe(diff_lf.group_by("global__kind").len())
    counts = {row["global__kind"]: int(row["len"]) for row in counts_df.iter_rows(named=True)}
    for key in ("changed", "left_only", "right_only", "unchanged"):
        counts.setdefault(key, 0)
    total = counts["changed"] + counts["left_only"] + counts["right_only"]

    top_changed = _top_changed_columns(
        diff_lf,
        compare_cols=compare_cols,
        dtype_map=dtype_map,
        rtol=rtol,
        atol=atol,
        limit=3,
        changed_rows=counts["changed"],
    )

    return {
        "keys": keys,
        "compare_cols": compare_cols,
        "left_extra": left_extra,
        "right_extra": right_extra,
        "rtol": rtol,
        "counts": counts,
        "total": total,
        "top_changed": top_changed,
        "left_path": left_path,
        "right_path": right_path,
        "row_mode": row_mode,
        "left_rows": left_rows,
        "right_rows": right_rows,
    }


def _diff_expr_for_column(
    name: str,
    *,
    dtype_map: pl.Schema,
    rtol: float | None,
    atol: float | None,
) -> pl.Expr:
    left_col = pl.col(f"{name}__left")
    right_col = pl.col(f"{name}__right")
    dtype = dtype_map.get(name)
    if (rtol is not None or atol is not None) and dtype is not None and _is_float_dtype(dtype):
        return _float_diff_expr(left_col, right_col, rtol=rtol, atol=atol)
    return left_col.eq_missing(right_col).not_()


def _top_changed_columns(
    diff_lf: pl.LazyFrame,
    *,
    compare_cols: list[str],
    dtype_map: pl.Schema,
    rtol: float | None,
    atol: float | None,
    limit: int,
    changed_rows: int,
) -> list[tuple[str, int, float]]:
    if not compare_cols or changed_rows <= 0 or limit <= 0:
        return []

    exprs = [
        _diff_expr_for_column(name, dtype_map=dtype_map, rtol=rtol, atol=atol).sum().alias(name)
        for name in compare_cols
    ]
    counts_df = collect_lazyframe(
        diff_lf.filter(pl.col("global__kind") == "changed").select(exprs).limit(1)
    )
    if counts_df.height == 0:
        return []

    row = counts_df.row(0)
    results: list[tuple[str, int, float]] = []
    for name, value in zip(counts_df.columns, row, strict=True):
        count = int(value or 0)
        if count <= 0:
            continue
        pct = count / changed_rows * 100 if changed_rows else 0.0
        results.append((name, count, pct))
    results.sort(key=lambda item: (-item[1], item[0]))
    return results[:limit]


def _row_count(lf: pl.LazyFrame) -> int:
    result = collect_lazyframe(lf.select(pl.len().alias("len")))
    if result.height == 0:
        return 0
    return int(result[0, 0])


__all__ = ["CompareResult", "compare_paths", "compare_lazyframes", "parse_name_list"]
