"""PandasBackend — translates Colnade expression trees and executes operations."""

from __future__ import annotations

import types as _types
from collections.abc import Iterator, Sequence
from typing import Any

import pandas as pd

from colnade.expr import (
    Agg,
    AliasedExpr,
    BinOp,
    ColumnRef,
    Expr,
    FunctionCall,
    ListOp,
    Literal,
    SortExpr,
    StructFieldAccess,
    UnaryOp,
)
from colnade.schema import Column, Schema, SchemaError
from colnade_pandas.conversion import map_colnade_dtype, map_pandas_dtype

# ---------------------------------------------------------------------------
# BinOp operator dispatch
# ---------------------------------------------------------------------------

_BINOP_MAP: dict[str, str] = {
    "+": "__add__",
    "-": "__sub__",
    "*": "__mul__",
    "/": "__truediv__",
    "%": "__mod__",
    ">": "__gt__",
    "<": "__lt__",
    ">=": "__ge__",
    "<=": "__le__",
    "==": "__eq__",
    "!=": "__ne__",
    "&": "__and__",
    "|": "__or__",
}

# ---------------------------------------------------------------------------
# Agg function name mapping (Colnade → Pandas)
# ---------------------------------------------------------------------------

_AGG_MAP: dict[str, str] = {
    "sum": "sum",
    "mean": "mean",
    "min": "min",
    "max": "max",
    "count": "count",
    "std": "std",
    "var": "var",
    "n_unique": "nunique",
    "first": "first",
    "last": "last",
}

# ---------------------------------------------------------------------------
# PandasBackend
# ---------------------------------------------------------------------------


class PandasBackend:
    """Colnade backend adapter for Pandas.

    Expression translation produces callables ``(df) -> Series | scalar``
    since Pandas has no standalone lazy expression API.
    """

    # --- Expression translation ---

    def translate_expr(self, expr: Expr[Any]) -> Any:
        """Recursively translate a Colnade AST node to a callable (df -> result)."""
        if isinstance(expr, AliasedExpr):
            inner_fn = self.translate_expr(expr.expr)
            target_name = expr.target.name
            # Return a tuple (fn, alias) so callers can unpack
            return (inner_fn, target_name)

        if isinstance(expr, ColumnRef):
            col_name = expr.column.name
            return lambda df, _cn=col_name: df[_cn]

        if isinstance(expr, Literal):
            val = expr.value
            return lambda df, _v=val: _v

        if isinstance(expr, BinOp):
            left_fn = self._ensure_callable(self.translate_expr(expr.left))
            right_fn = self._ensure_callable(self.translate_expr(expr.right))
            method = _BINOP_MAP.get(expr.op)
            if method is None:
                msg = f"Unsupported BinOp operator: {expr.op}"
                raise ValueError(msg)
            return lambda df, _l=left_fn, _r=right_fn, _m=method: getattr(_l(df), _m)(_r(df))

        if isinstance(expr, UnaryOp):
            operand_fn = self._ensure_callable(self.translate_expr(expr.operand))
            if expr.op == "-":
                return lambda df, _o=operand_fn: -_o(df)
            if expr.op == "~":
                return lambda df, _o=operand_fn: ~_o(df)
            if expr.op == "is_null":
                return lambda df, _o=operand_fn: _o(df).isna()
            if expr.op == "is_not_null":
                return lambda df, _o=operand_fn: _o(df).notna()
            if expr.op == "is_nan":
                return lambda df, _o=operand_fn: _o(df).isna()
            msg = f"Unsupported UnaryOp: {expr.op}"
            raise ValueError(msg)

        if isinstance(expr, Agg):
            source_fn = self._ensure_callable(self.translate_expr(expr.source))
            agg_name = _AGG_MAP.get(expr.agg_type)
            if agg_name is None:
                msg = f"Unsupported aggregation: {expr.agg_type}"
                raise ValueError(msg)
            return (source_fn, agg_name)

        if isinstance(expr, FunctionCall):
            return self._translate_function_call(expr)

        if isinstance(expr, StructFieldAccess):
            struct_fn = self._ensure_callable(self.translate_expr(expr.struct_expr))
            field_name = expr.field.name
            return lambda df, _s=struct_fn, _f=field_name: _s(df).apply(lambda x: x.get(_f))

        if isinstance(expr, ListOp):
            return self._translate_list_op(expr)

        msg = f"Unsupported expression type: {type(expr).__name__}"
        raise TypeError(msg)

    def _ensure_callable(self, translated: Any) -> Any:
        """Unwrap (fn, alias) tuples from AliasedExpr to get the callable."""
        if isinstance(translated, tuple):
            return translated[0]
        return translated

    def _translate_function_call(self, expr: FunctionCall[Any]) -> Any:
        """Translate a FunctionCall node to a Pandas callable."""
        name = expr.name

        # String methods
        if name == "str_contains":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            pattern = expr.args[1]
            return lambda df, _s=source_fn, _p=pattern: _s(df).str.contains(_p, regex=False)
        if name == "str_starts_with":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            pattern = expr.args[1]
            return lambda df, _s=source_fn, _p=pattern: _s(df).str.startswith(_p)
        if name == "str_ends_with":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            pattern = expr.args[1]
            return lambda df, _s=source_fn, _p=pattern: _s(df).str.endswith(_p)
        if name == "str_len":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            return lambda df, _s=source_fn: _s(df).str.len()
        if name == "str_to_lowercase":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            return lambda df, _s=source_fn: _s(df).str.lower()
        if name == "str_to_uppercase":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            return lambda df, _s=source_fn: _s(df).str.upper()
        if name == "str_strip":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            return lambda df, _s=source_fn: _s(df).str.strip()
        if name == "str_replace":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            old, new = expr.args[1], expr.args[2]
            return lambda df, _s=source_fn, _o=old, _n=new: _s(df).str.replace(_o, _n, regex=False)

        # Temporal methods
        if name == "dt_year":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            return lambda df, _s=source_fn: _s(df).dt.year
        if name == "dt_month":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            return lambda df, _s=source_fn: _s(df).dt.month
        if name == "dt_day":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            return lambda df, _s=source_fn: _s(df).dt.day
        if name == "dt_hour":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            return lambda df, _s=source_fn: _s(df).dt.hour
        if name == "dt_minute":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            return lambda df, _s=source_fn: _s(df).dt.minute
        if name == "dt_second":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            return lambda df, _s=source_fn: _s(df).dt.second
        if name == "dt_truncate":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            unit = expr.args[1]
            # Pandas uses floor with offset aliases
            return lambda df, _s=source_fn, _u=unit: _s(df).dt.floor(_u)

        # Null/NaN handling
        if name == "fill_null":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            fill_fn = self._ensure_callable(self.translate_expr(expr.args[1]))
            return lambda df, _s=source_fn, _f=fill_fn: _s(df).fillna(_f(df))
        if name == "fill_nan":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            fill_fn = self._ensure_callable(self.translate_expr(expr.args[1]))
            return lambda df, _s=source_fn, _f=fill_fn: _s(df).fillna(_f(df))
        if name == "assert_non_null":
            return self._ensure_callable(self.translate_expr(expr.args[0]))

        # Cast
        if name == "cast":
            source_fn = self._ensure_callable(self.translate_expr(expr.args[0]))
            target_dtype = map_colnade_dtype(expr.kwargs["dtype"])
            return lambda df, _s=source_fn, _t=target_dtype: _s(df).astype(_t)

        # Window function (over)
        if name == "over":
            raw_source = self.translate_expr(expr.args[0])
            partition_fns = [self._ensure_callable(self.translate_expr(a)) for a in expr.args[1:]]
            if isinstance(raw_source, tuple):
                col_fn, agg_name = raw_source
                return lambda df, _c=col_fn, _a=agg_name, _p=partition_fns: df.groupby(
                    [p(df).name for p in _p]
                )[_c(df).name].transform(_a)
            source_fn = raw_source
            return lambda df, _s=source_fn, _p=partition_fns: df.groupby([p(df).name for p in _p])[
                _s(df).name
            ].transform(lambda x: x)

        msg = f"Unsupported FunctionCall: {name}"
        raise ValueError(msg)

    def _translate_list_op(self, expr: ListOp[Any]) -> Any:
        """Translate a ListOp node to a Pandas callable."""
        list_fn = self._ensure_callable(self.translate_expr(expr.list_expr))
        op = expr.op

        if op == "len":
            return lambda df, _l=list_fn: _l(df).apply(len)
        if op == "get":
            idx = expr.args[0]
            return lambda df, _l=list_fn, _i=idx: _l(df).apply(lambda x: x[_i])
        if op == "contains":
            val = expr.args[0]
            return lambda df, _l=list_fn, _v=val: _l(df).apply(lambda x: _v in x)
        if op == "sum":
            return lambda df, _l=list_fn: _l(df).apply(sum)
        if op == "mean":
            return lambda df, _l=list_fn: _l(df).apply(lambda x: sum(x) / len(x) if x else None)
        if op == "min":
            return lambda df, _l=list_fn: _l(df).apply(min)
        if op == "max":
            return lambda df, _l=list_fn: _l(df).apply(max)

        msg = f"Unsupported ListOp: {op}"
        raise ValueError(msg)

    # --- Execution methods ---

    def filter(self, source: Any, predicate: Expr[Any]) -> Any:
        pred_fn = self._ensure_callable(self.translate_expr(predicate))
        mask = pred_fn(source)
        return source.loc[mask].reset_index(drop=True)

    def sort(
        self,
        source: Any,
        by: Sequence[Column[Any] | SortExpr],
        descending: bool,
    ) -> Any:
        col_names: list[str] = []
        ascending: list[bool] = []
        for item in by:
            if isinstance(item, SortExpr):
                col_names.append(item.expr.column.name)
                ascending.append(not item.descending)
            else:
                col_names.append(item.name)
                ascending.append(not descending)
        return source.sort_values(by=col_names, ascending=ascending).reset_index(drop=True)

    def limit(self, source: Any, n: int) -> Any:
        return source.head(n).reset_index(drop=True)

    def head(self, source: Any, n: int) -> Any:
        return source.head(n).reset_index(drop=True)

    def tail(self, source: Any, n: int) -> Any:
        return source.tail(n).reset_index(drop=True)

    def sample(self, source: Any, n: int) -> Any:
        return source.sample(n).reset_index(drop=True)

    def unique(self, source: Any, columns: Sequence[Column[Any]]) -> Any:
        subset = [c.name for c in columns] or None
        return source.drop_duplicates(subset=subset).reset_index(drop=True)

    def drop_nulls(self, source: Any, columns: Sequence[Column[Any]]) -> Any:
        return source.dropna(subset=[c.name for c in columns]).reset_index(drop=True)

    def with_columns(self, source: Any, exprs: Sequence[AliasedExpr[Any] | Expr[Any]]) -> Any:
        result = source.copy()
        for expr in exprs:
            translated = self.translate_expr(expr)
            if isinstance(translated, tuple):
                fn, alias = translated
                result[alias] = fn(result)
            else:
                # Non-aliased expression — shouldn't normally happen in with_columns
                msg = "with_columns requires aliased expressions"
                raise ValueError(msg)
        return result

    def select(self, source: Any, columns: Sequence[Column[Any]]) -> Any:
        return source[[c.name for c in columns]].reset_index(drop=True)

    def group_by_agg(
        self,
        source: Any,
        keys: Sequence[Column[Any]],
        aggs: Sequence[AliasedExpr[Any]],
    ) -> Any:
        key_names = [k.name for k in keys]
        agg_dict: dict[str, tuple[str, str]] = {}
        rename_map: dict[str, str] = {}

        for agg_expr in aggs:
            translated = self.translate_expr(agg_expr)
            if not isinstance(translated, tuple):
                msg = "group_by_agg requires aliased aggregation expressions"
                raise ValueError(msg)

            inner, alias = translated
            # inner is (source_fn, agg_name) from Agg translation
            if isinstance(inner, tuple):
                source_fn, agg_name = inner
                # Get the column name from the source function
                # This works for simple ColumnRef-based aggregations
                col_name = self._extract_col_name(source_fn, source)
                agg_dict[col_name] = agg_name
                rename_map[col_name] = alias
            else:
                msg = "group_by_agg requires aggregation expressions (e.g., .sum(), .mean())"
                raise ValueError(msg)

        grouped = source.groupby(key_names, as_index=False).agg(agg_dict)
        # Rename aggregated columns to their aliases
        for old_name, new_name in rename_map.items():
            if old_name != new_name and old_name not in key_names:
                grouped = grouped.rename(columns={old_name: new_name})
        return grouped

    def agg(self, source: Any, aggs: Sequence[AliasedExpr[Any]]) -> Any:
        result: dict[str, Any] = {}
        for agg_expr in aggs:
            translated = self.translate_expr(agg_expr)
            inner, alias = translated
            source_fn, agg_name = inner
            col_name = self._extract_col_name(source_fn, source)
            result[alias] = getattr(source[col_name], agg_name)()
        return pd.DataFrame([result])

    def _extract_col_name(self, fn: Any, df: pd.DataFrame) -> str:
        """Extract column name from a translated ColumnRef function."""
        series = fn(df)
        if hasattr(series, "name"):
            return series.name
        msg = "Cannot extract column name from expression"
        raise ValueError(msg)

    def join(self, left: Any, right: Any, on: Any, how: str) -> Any:
        return left.merge(right, left_on=on.left.name, right_on=on.right.name, how=how)

    def cast_schema(self, source: Any, column_mapping: dict[str, str]) -> Any:
        # column_mapping is {target_name: source_name}
        # Build the reverse mapping for rename, then select target columns
        rename_map = {src: tgt for tgt, src in column_mapping.items()}
        result = source.rename(columns=rename_map)
        return result[list(column_mapping.keys())].reset_index(drop=True)

    def lazy(self, source: Any) -> Any:
        # Pandas has no lazy mode — passthrough
        return source

    def collect(self, source: Any) -> Any:
        # Pandas has no lazy mode — passthrough
        return source

    @staticmethod
    def _dtypes_compatible(actual: Any, expected: Any) -> bool:
        """Compare dtypes accounting for NumPy vs Pandas extension type differences.

        Pandas may return NumPy dtypes (e.g. numpy.uint64) for data read from
        parquet/CSV while Colnade maps to extension types (e.g. pd.UInt64Dtype()).
        """
        if actual == expected:
            return True
        if isinstance(actual, pd.StringDtype) and isinstance(expected, pd.StringDtype):
            return True
        import numpy as np

        if isinstance(actual, np.dtype) and hasattr(expected, "numpy_dtype"):
            return actual == expected.numpy_dtype
        return False

    def validate_schema(self, source: Any, schema: type[Schema]) -> None:
        """Validate that a Pandas DataFrame matches the schema."""
        expected_columns = schema._columns
        actual_names = set(source.columns)

        missing = [n for n in expected_columns if n not in actual_names]
        type_mismatches: dict[str, tuple[str, str]] = {}

        for col_name, col in expected_columns.items():
            if col_name not in actual_names:
                continue
            actual_pd_dtype = source[col_name].dtype
            expected_pd_dtype = map_colnade_dtype(col.dtype)
            if not self._dtypes_compatible(actual_pd_dtype, expected_pd_dtype):
                try:
                    actual_colnade = map_pandas_dtype(actual_pd_dtype).__name__
                except TypeError:
                    actual_colnade = str(actual_pd_dtype)
                expected_name = (
                    col.dtype.__name__ if hasattr(col.dtype, "__name__") else str(col.dtype)
                )
                type_mismatches[col_name] = (expected_name, actual_colnade)

        # Check nullability — non-nullable columns should have no nulls
        null_violations: list[str] = []
        for col_name, col in expected_columns.items():
            if col_name not in actual_names:
                continue
            if isinstance(col.dtype, _types.UnionType):
                continue
            if source[col_name].isna().any():
                null_violations.append(col_name)

        if missing or type_mismatches or null_violations:
            raise SchemaError(
                missing_columns=missing if missing else None,
                type_mismatches=type_mismatches if type_mismatches else None,
                null_violations=null_violations if null_violations else None,
            )

    def validate_field_constraints(self, source: Any, schema: type[Schema]) -> None:
        """Validate value-level constraints (Field(), @schema_check) on data."""
        from colnade.constraints import ValueViolation, get_column_constraints, get_schema_checks

        constraints = get_column_constraints(schema)
        checks = get_schema_checks(schema)
        if not constraints and not checks:
            return

        violations: list[ValueViolation] = []

        for col_name, field_info in constraints.items():
            if col_name not in source.columns:
                continue
            series = source[col_name].dropna()

            if field_info.ge is not None:
                mask = series < field_info.ge
                count = int(mask.sum())
                if count > 0:
                    samples = series[mask].head(5).tolist()
                    violations.append(
                        ValueViolation(
                            column=col_name,
                            constraint=f"ge={field_info.ge!r}",
                            got_count=count,
                            sample_values=samples,
                        )
                    )

            if field_info.gt is not None:
                mask = series <= field_info.gt
                count = int(mask.sum())
                if count > 0:
                    samples = series[mask].head(5).tolist()
                    violations.append(
                        ValueViolation(
                            column=col_name,
                            constraint=f"gt={field_info.gt!r}",
                            got_count=count,
                            sample_values=samples,
                        )
                    )

            if field_info.le is not None:
                mask = series > field_info.le
                count = int(mask.sum())
                if count > 0:
                    samples = series[mask].head(5).tolist()
                    violations.append(
                        ValueViolation(
                            column=col_name,
                            constraint=f"le={field_info.le!r}",
                            got_count=count,
                            sample_values=samples,
                        )
                    )

            if field_info.lt is not None:
                mask = series >= field_info.lt
                count = int(mask.sum())
                if count > 0:
                    samples = series[mask].head(5).tolist()
                    violations.append(
                        ValueViolation(
                            column=col_name,
                            constraint=f"lt={field_info.lt!r}",
                            got_count=count,
                            sample_values=samples,
                        )
                    )

            if field_info.min_length is not None:
                lengths = series.str.len()
                mask = lengths < field_info.min_length
                count = int(mask.sum())
                if count > 0:
                    samples = series[mask].head(5).tolist()
                    violations.append(
                        ValueViolation(
                            column=col_name,
                            constraint=f"min_length={field_info.min_length}",
                            got_count=count,
                            sample_values=samples,
                        )
                    )

            if field_info.max_length is not None:
                lengths = series.str.len()
                mask = lengths > field_info.max_length
                count = int(mask.sum())
                if count > 0:
                    samples = series[mask].head(5).tolist()
                    violations.append(
                        ValueViolation(
                            column=col_name,
                            constraint=f"max_length={field_info.max_length}",
                            got_count=count,
                            sample_values=samples,
                        )
                    )

            if field_info.pattern is not None:
                matches = series.str.contains(field_info.pattern, regex=True, na=False)
                mask = ~matches
                count = int(mask.sum())
                if count > 0:
                    samples = series[mask].head(5).tolist()
                    violations.append(
                        ValueViolation(
                            column=col_name,
                            constraint=f"pattern={field_info.pattern!r}",
                            got_count=count,
                            sample_values=samples,
                        )
                    )

            if field_info.unique:
                dup_mask = series.duplicated(keep=False)
                count = int(dup_mask.sum())
                if count > 0:
                    samples = series[dup_mask].unique().tolist()[:5]
                    violations.append(
                        ValueViolation(
                            column=col_name,
                            constraint="unique",
                            got_count=count,
                            sample_values=samples,
                        )
                    )

            if field_info.isin is not None:
                mask = ~series.isin(field_info.isin)
                count = int(mask.sum())
                if count > 0:
                    samples = series[mask].unique().tolist()[:5]
                    violations.append(
                        ValueViolation(
                            column=col_name,
                            constraint=f"isin={list(field_info.isin)!r}",
                            got_count=count,
                            sample_values=samples,
                        )
                    )

        # Cross-column schema checks
        for check in checks:
            expr = check.fn(schema)
            pd_mask = self.translate_expr(expr)
            violation_count = int((~pd_mask(source)).sum())
            if violation_count > 0:
                violations.append(
                    ValueViolation(
                        column="(cross-column)",
                        constraint=f"schema_check:{check.name}",
                        got_count=violation_count,
                        sample_values=[],
                    )
                )

        if violations:
            raise SchemaError(value_violations=violations)

    # --- Introspection ---

    def row_count(self, source: Any) -> int:
        return len(source)

    def iter_row_dicts(self, source: Any) -> Iterator[dict[str, Any]]:
        return source.to_dict(orient="records")

    # --- Arrow boundary ---

    def to_arrow_batches(
        self,
        source: Any,
        batch_size: int | None,
    ) -> Iterator[Any]:
        """Convert a Pandas DataFrame to an iterator of Arrow RecordBatches."""
        import pyarrow as pa

        table: pa.Table = pa.Table.from_pandas(source)
        if batch_size is not None:
            yield from table.to_batches(max_chunksize=batch_size)
        else:
            yield from table.to_batches()

    def from_arrow_batches(
        self,
        batches: Iterator[Any],
        schema: type[Any],
    ) -> Any:
        """Reconstruct a Pandas DataFrame from Arrow RecordBatches."""
        import pyarrow as pa

        batch_list = list(batches)
        if not batch_list:
            return pd.DataFrame()
        table = pa.Table.from_batches(batch_list)
        return table.to_pandas()

    def from_dict(
        self,
        data: dict[str, Sequence[Any]],
        schema: type[Schema],
    ) -> Any:
        """Create a Pandas DataFrame from a columnar dict with schema-driven dtypes."""
        from colnade_pandas.conversion import map_colnade_dtype

        pd_schema = {name: map_colnade_dtype(col.dtype) for name, col in schema._columns.items()}
        df = pd.DataFrame(data)
        return df.astype(pd_schema)
