from __future__ import annotations

import operator
import warnings
from functools import reduce
from itertools import chain
from typing import TYPE_CHECKING, Any, Literal, Protocol, cast, overload

import pandas as pd

from narwhals._compliant import EagerNamespace
from narwhals._expression_parsing import (
    combine_alias_output_names,
    combine_evaluate_output_names,
)
from narwhals._pandas_like.dataframe import PandasLikeDataFrame
from narwhals._pandas_like.expr import PandasLikeExpr
from narwhals._pandas_like.selectors import PandasSelectorNamespace
from narwhals._pandas_like.series import PandasLikeSeries
from narwhals._pandas_like.typing import NativeDataFrameT, NativeSeriesT
from narwhals._pandas_like.utils import (
    import_array_module,
    is_dtype_pyarrow,
    is_non_nullable_boolean,
    set_index,
)
from narwhals._utils import validate_separators

if TYPE_CHECKING:
    from collections.abc import Iterable, Sequence
    from typing import TypeAlias

    from narwhals._utils import Implementation, Version
    from narwhals.typing import (
        CorrelationMethod,
        IntoDType,
        NormalizedPath,
        PythonLiteral,
    )


Incomplete: TypeAlias = Any
"""Escape hatch, but leaving a trace that this isn't ideal."""


_Vertical: TypeAlias = Literal[0]
_Horizontal: TypeAlias = Literal[1]
Axis: TypeAlias = Literal[_Vertical, _Horizontal]

VERTICAL: _Vertical = 0
HORIZONTAL: _Horizontal = 1


class PandasLikeNamespace(
    EagerNamespace[
        PandasLikeDataFrame,
        PandasLikeSeries,
        PandasLikeExpr,
        NativeDataFrameT,
        NativeSeriesT,
    ]
):
    @property
    def _dataframe(self) -> type[PandasLikeDataFrame]:
        return PandasLikeDataFrame

    @property
    def _expr(self) -> type[PandasLikeExpr]:
        return PandasLikeExpr

    @property
    def _series(self) -> type[PandasLikeSeries]:
        return PandasLikeSeries

    @property
    def selectors(self) -> PandasSelectorNamespace:
        return PandasSelectorNamespace.from_namespace(self)

    def __init__(self, implementation: Implementation, version: Version) -> None:
        self._implementation = implementation
        self._version = version

    def read_csv(
        self, source: NormalizedPath, *, separator: str = ",", **kwds: Any
    ) -> PandasLikeDataFrame:
        validate_separators(separator, ("sep",), kwds)
        ns = self._implementation.to_native_namespace()
        native = ns.read_csv(source, sep=separator, **kwds)
        return self._dataframe.from_native(native, context=self)

    def read_parquet(self, source: NormalizedPath, **kwds: Any) -> PandasLikeDataFrame:
        ns = self._implementation.to_native_namespace()
        return self._dataframe.from_native(ns.read_parquet(source, **kwds), context=self)

    def coalesce(self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
            series = (s for _expr in exprs for s in _expr(df))
            return [
                reduce(lambda x, y: x.fill_null(y, strategy=None, limit=None), series)
            ]

        return self._expr._from_callable(
            func=func,
            evaluate_output_names=combine_evaluate_output_names(*exprs),
            alias_output_names=combine_alias_output_names(*exprs),
            context=self,
        )

    def lit(self, value: PythonLiteral, dtype: IntoDType | None) -> PandasLikeExpr:
        def _lit_pandas_series(df: PandasLikeDataFrame) -> PandasLikeSeries:
            if isinstance(value, (list, tuple, dict)):
                try:
                    import pandas as pd  # ignore-banned-import
                    import pyarrow as pa  # ignore-banned-import
                except ImportError as exc:  # pragma: no cover
                    msg = (
                        "Nested structures require pyarrow to be installed for pandas backend. "
                        "Please install pyarrow: pip install pyarrow"
                    )
                    raise ImportError(msg) from exc

                from narwhals._arrow.utils import (
                    narwhals_to_native_dtype as _to_arrow_dtype,
                )

                array_value = list(value) if isinstance(value, tuple) else value
                pa_dtype = _to_arrow_dtype(dtype, self._version) if dtype else None
                pa_array = pa.array([array_value], type=pa_dtype)  # type: ignore[arg-type, list-item]

                # Use ArrowExtensionArray to avoid pandas unpacking the nested structure
                ns = self._implementation.to_native_namespace()
                pandas_series_native = ns.Series(
                    pd.arrays.ArrowExtensionArray(pa_array),  # type: ignore[attr-defined]
                    name="literal",
                    index=df._native_frame.index[0:1],
                )

                return self._series.from_native(pandas_series_native, context=self)

            pandas_like_series = self._series.from_iterable(
                data=[value],
                name="literal",
                index=df._native_frame.index[0:1],
                context=self,
            )
            if dtype:
                return pandas_like_series.cast(dtype)
            return pandas_like_series

        return PandasLikeExpr(
            lambda df: [_lit_pandas_series(df)],
            evaluate_output_names=lambda _df: ["literal"],
            alias_output_names=None,
            implementation=self._implementation,
            version=self._version,
        )

    def len(self) -> PandasLikeExpr:
        return PandasLikeExpr(
            lambda df: [
                self._series.from_iterable(
                    [len(df._native_frame)], name="len", index=[0], context=self
                )
            ],
            evaluate_output_names=lambda _df: ["len"],
            alias_output_names=None,
            implementation=self._implementation,
            version=self._version,
        )

    # --- horizontal ---
    def sum_horizontal(self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
            it = chain.from_iterable(expr(df) for expr in exprs)
            native_series = (s.fill_null(0, None, None) for s in it)
            return [reduce(operator.add, native_series)]

        return self._expr._from_callable(
            func=func,
            evaluate_output_names=combine_evaluate_output_names(*exprs),
            alias_output_names=combine_alias_output_names(*exprs),
            context=self,
        )

    def all_horizontal(
        self, *exprs: PandasLikeExpr, ignore_nulls: bool
    ) -> PandasLikeExpr:
        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
            series = [s for _expr in exprs for s in _expr(df)]
            if not ignore_nulls and any(
                s.native.dtype == "object" and s.is_null().any() for s in series
            ):
                # classical NumPy boolean columns don't support missing values, so
                # only do the full scan with `is_null` if we have `object` dtype.
                msg = "Cannot use `ignore_nulls=False` in `all_horizontal` for non-nullable NumPy-backed pandas Series when nulls are present."
                raise ValueError(msg)
            it = (
                (
                    # NumPy-backed 'bool' dtype can't contain nulls so doesn't need filling.
                    s if is_non_nullable_boolean(s) else s.fill_null(True, None, None)
                    for s in series
                )
                if ignore_nulls
                else iter(series)
            )
            return [reduce(operator.and_, it)]

        return self._expr._from_callable(
            func=func,
            evaluate_output_names=combine_evaluate_output_names(*exprs),
            alias_output_names=combine_alias_output_names(*exprs),
            context=self,
        )

    def any_horizontal(
        self, *exprs: PandasLikeExpr, ignore_nulls: bool
    ) -> PandasLikeExpr:
        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
            series = [s for _expr in exprs for s in _expr(df)]
            if not ignore_nulls and any(
                s.native.dtype == "object" and s.is_null().any() for s in series
            ):
                # classical NumPy boolean columns don't support missing values, so
                # only do the full scan with `is_null` if we have `object` dtype.
                msg = "Cannot use `ignore_nulls=False` in `any_horizontal` for non-nullable NumPy-backed pandas Series when nulls are present."
                raise ValueError(msg)
            it = (
                (
                    # NumPy-backed 'bool' dtype can't contain nulls so doesn't need filling.
                    s if is_non_nullable_boolean(s) else s.fill_null(False, None, None)
                    for s in series
                )
                if ignore_nulls
                else iter(series)
            )
            return [reduce(operator.or_, it)]

        return self._expr._from_callable(
            func=func,
            evaluate_output_names=combine_evaluate_output_names(*exprs),
            alias_output_names=combine_alias_output_names(*exprs),
            context=self,
        )

    def mean_horizontal(self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
            expr_results = [s for _expr in exprs for s in _expr(df)]
            series = (s.fill_null(0, strategy=None, limit=None) for s in expr_results)
            non_na = (1 - s.is_null() for s in expr_results)
            return [reduce(operator.add, series) / reduce(operator.add, non_na)]

        return self._expr._from_callable(
            func=func,
            evaluate_output_names=combine_evaluate_output_names(*exprs),
            alias_output_names=combine_alias_output_names(*exprs),
            context=self,
        )

    def min_horizontal(self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
            series = list(chain.from_iterable(expr(df) for expr in exprs))
            return [
                PandasLikeSeries(
                    self.concat(
                        (s.to_frame() for s in series), how="horizontal"
                    )._native_frame.min(axis=1),
                    implementation=self._implementation,
                    version=self._version,
                ).alias(series[0].name)
            ]

        return self._expr._from_callable(
            func=func,
            evaluate_output_names=combine_evaluate_output_names(*exprs),
            alias_output_names=combine_alias_output_names(*exprs),
            context=self,
        )

    def max_horizontal(self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
            series = list(chain.from_iterable(expr(df) for expr in exprs))
            return [
                PandasLikeSeries(
                    self.concat(
                        (s.to_frame() for s in series), how="horizontal"
                    ).native.max(axis=1),
                    implementation=self._implementation,
                    version=self._version,
                ).alias(series[0].name)
            ]

        return self._expr._from_callable(
            func=func,
            evaluate_output_names=combine_evaluate_output_names(*exprs),
            alias_output_names=combine_alias_output_names(*exprs),
            context=self,
        )

    @property
    def _concat(self) -> _NativeConcat[NativeDataFrameT, NativeSeriesT]:
        """Concatenate pandas objects along a particular axis.

        Return the **native** equivalent of `pd.concat`.
        """
        return self._implementation.to_native_namespace().concat

    def _concat_diagonal(self, dfs: Sequence[NativeDataFrameT], /) -> NativeDataFrameT:
        if self._implementation.is_pandas() and self._backend_version < (3,):
            return self._concat(dfs, axis=VERTICAL, copy=False)
        return self._concat(dfs, axis=VERTICAL)

    def _concat_by_index(
        self, dfs: Sequence[NativeDataFrameT | NativeSeriesT], /
    ) -> NativeDataFrameT:
        """Concatenate horizontally, aligning inputs by their (label-based) index.

        Use this only when inputs are already known to share a meaningful,
        aligned index (e.g. multiple aggregations from the same `group_by`).
        For the general case, where inputs' indices may be unrelated, use
        `_concat_horizontal` instead.
        """
        if self._implementation.is_cudf():
            with warnings.catch_warnings():
                warnings.filterwarnings(
                    "ignore",
                    message="The behavior of array concatenation with empty entries is deprecated",
                    category=FutureWarning,
                )
                return self._concat(dfs, axis=HORIZONTAL)
        elif self._implementation.is_pandas() and self._backend_version < (3,):
            return self._concat(dfs, axis=HORIZONTAL, copy=False)
        return self._concat(dfs, axis=HORIZONTAL)

    def _concat_horizontal(self, dfs: Sequence[NativeDataFrameT], /) -> NativeDataFrameT:
        # Follow the left-hand-rule documented in `docs/concepts/pandas_index.md`.
        target_index = dfs[0].index
        arange = import_array_module(self._implementation).arange
        impl = self._implementation

        def reset(df: NativeDataFrameT) -> NativeDataFrameT:
            # Value of type variable "NativeNDFrameT" of "set_index" cannot be "NativeDataFrameT"
            return set_index(  # type: ignore[type-var]
                df, arange(len(df)), implementation=impl
            )

        concatenated = self._concat_by_index([reset(df) for df in dfs])
        # Value of type variable "NativeNDFrameT" of "set_index" cannot be "NativeDataFrameT"
        return set_index(  # type: ignore[type-var]
            concatenated, target_index, implementation=impl
        )

    def _concat_vertical(self, dfs: Sequence[NativeDataFrameT], /) -> NativeDataFrameT:
        cols_0 = dfs[0].columns
        for i, df in enumerate(dfs[1:], start=1):
            cols_current = df.columns
            if not (
                (len(cols_current) == len(cols_0)) and (cols_current == cols_0).all()
            ):
                msg = (
                    "unable to vstack, column names don't match:\n"
                    f"   - dataframe 0: {cols_0.to_list()}\n"
                    f"   - dataframe {i}: {cols_current.to_list()}\n"
                )
                raise TypeError(msg)
        if self._implementation.is_pandas() and self._backend_version < (3,):
            return self._concat(dfs, axis=VERTICAL, copy=False)
        return self._concat(dfs, axis=VERTICAL)

    def concat_str(
        self, *exprs: PandasLikeExpr, separator: str, ignore_nulls: bool
    ) -> PandasLikeExpr:
        string = self._version.dtypes.String()

        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
            expr_results = [s for _expr in exprs for s in _expr(df)]
            series = [s.cast(string) for s in expr_results]
            null_mask = [s.is_null() for s in expr_results]

            if not ignore_nulls:
                null_mask_result = reduce(operator.or_, null_mask)
                result = reduce(lambda x, y: x + separator + y, series).zip_with(
                    ~null_mask_result, None
                )
            else:
                # NOTE: Trying to help `mypy` later
                # error: Cannot determine type of "values"  [has-type]
                values: list[PandasLikeSeries]
                init_value, *values = (
                    s.zip_with(~nm, "") for s, nm in zip(series, null_mask, strict=True)
                )
                sep_array = init_value._with_native(
                    init_value.__native_namespace__().Series(
                        separator,
                        name="sep",
                        index=init_value.native.index,
                        dtype=init_value.native.dtype,
                    )
                )
                separators = (sep_array.zip_with(~nm, "") for nm in null_mask[:-1])
                result = reduce(
                    operator.add,
                    (s + v for s, v in zip(separators, values, strict=True)),
                    init_value,
                )

            return [result]

        return self._expr._from_callable(
            func=func,
            evaluate_output_names=combine_evaluate_output_names(*exprs),
            alias_output_names=combine_alias_output_names(*exprs),
            context=self,
        )

    def struct(self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
            try:
                import pandas as pd  # ignore-banned-import
                import pyarrow as pa  # ignore-banned-import
                import pyarrow.compute as pc  # ignore-banned-import
            except ImportError as exc:  # pragma: no cover
                msg = (
                    "struct requires pyarrow to be installed for pandas backend. "
                    "Please install pyarrow: pip install pyarrow"
                )
                raise ImportError(msg) from exc

            align = self._series._align_full_broadcast
            series = align(*chain.from_iterable(expr(df) for expr in exprs))
            name = series[0].name
            struct_array = pc.make_struct(
                *(pa.array(s.native, from_pandas=True) for s in series),
                field_names=[s.name for s in series],
            )
            result = pa.chunked_array([struct_array])

            version = self._version
            impl = self._implementation
            ns = impl.to_native_namespace()

            result_native = ns.Series(
                pd.arrays.ArrowExtensionArray(result),  # type: ignore[attr-defined]
                name=name,
                index=series[0].native.index,
            )
            return [self._series(result_native, implementation=impl, version=version)]

        return self._expr._from_callable(
            func=func,
            evaluate_output_names=combine_evaluate_output_names(*exprs),
            alias_output_names=combine_alias_output_names(*exprs),
            context=self,
        )

    def list(self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
            try:
                import pyarrow as pa  # ignore-banned-import
            except ImportError as exc:  # pragma: no cover
                msg = (
                    "list requires pyarrow to be installed for pandas backend. "
                    "Please install pyarrow: pip install pyarrow"
                )
                raise ImportError(msg) from exc

            from narwhals._arrow.utils import build_list_array

            align = self._series._align_full_broadcast
            series = align(*chain.from_iterable(expr(df) for expr in exprs))
            arrays = [
                (
                    s.native.array._pa_array.combine_chunks()
                    if is_dtype_pyarrow(s.native.dtype)
                    else pa.array(s.native, from_pandas=True)
                )
                for s in series
            ]
            result = build_list_array(arrays)

            impl = self._implementation
            result_native = impl.to_native_namespace().Series(
                pd.arrays.ArrowExtensionArray(result),  # type: ignore[attr-defined]
                name=series[0].name,
                index=series[0].native.index,
            )
            return [
                self._series(result_native, implementation=impl, version=self._version)
            ]

        return self._expr._from_callable(
            func=func,
            evaluate_output_names=combine_evaluate_output_names(*exprs),
            alias_output_names=combine_alias_output_names(*exprs),
            context=self,
        )

    def _if_then_else(
        self,
        when: NativeSeriesT,
        then: NativeSeriesT,
        otherwise: NativeSeriesT | None = None,
    ) -> NativeSeriesT:
        where: Incomplete = then.where
        return where(when) if otherwise is None else where(when, otherwise)

    def corr(
        self, a: PandasLikeExpr, b: PandasLikeExpr, *, method: CorrelationMethod
    ) -> PandasLikeExpr:
        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
            a_series = df._evaluate_single_output_expr(a)
            b_series = df._evaluate_single_output_expr(b)
            aligned = cast(
                "list[NativeSeriesT]",
                [df._extract_comparand(a_series), df._extract_comparand(b_series)],
            )
            _df = self._concat_by_index(aligned)
            corr = _df.corr(method=method).iloc[0, [1]]  # type: ignore[union-attr]
            return [
                PandasLikeSeries(
                    corr, implementation=self._implementation, version=self._version
                )
            ]

        return self._expr._from_callable(
            func=func,
            evaluate_output_names=combine_evaluate_output_names(a, b),
            alias_output_names=combine_alias_output_names(a, b),
            context=self,
        )

    def cov(self, a: PandasLikeExpr, b: PandasLikeExpr, *, ddof: int) -> PandasLikeExpr:
        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
            a_series = df._evaluate_single_output_expr(a)
            b_series = df._evaluate_single_output_expr(b)
            aligned = cast(
                "list[NativeSeriesT]",
                [df._extract_comparand(a_series), df._extract_comparand(b_series)],
            )
            _df = cast("pd.DataFrame", self._concat_by_index(aligned))
            n = _df.count(axis=1).eq(2).sum()
            if ddof == 0 and n == 1:
                value = 0.0
            elif n - ddof <= 0:
                value = float("nan")
            else:
                value = _df.cov(ddof=ddof).iloc[0, 1]
            return [
                PandasLikeSeries.from_iterable([value], context=self, name=a_series.name)
            ]

        return self._expr._from_callable(
            func=func,
            evaluate_output_names=combine_evaluate_output_names(a, b),
            alias_output_names=combine_alias_output_names(a, b),
            context=self,
        )


class _NativeConcat(Protocol[NativeDataFrameT, NativeSeriesT]):
    @overload
    def __call__(
        self,
        objs: Iterable[NativeDataFrameT],
        *,
        axis: _Vertical,
        copy: bool | None = ...,
    ) -> NativeDataFrameT: ...
    @overload
    def __call__(
        self, objs: Iterable[NativeSeriesT], *, axis: _Vertical, copy: bool | None = ...
    ) -> NativeSeriesT: ...
    @overload
    def __call__(
        self,
        objs: Iterable[NativeDataFrameT | NativeSeriesT],
        *,
        axis: _Horizontal,
        copy: bool | None = ...,
    ) -> NativeDataFrameT: ...
    @overload
    def __call__(
        self,
        objs: Iterable[NativeDataFrameT | NativeSeriesT],
        *,
        axis: Axis,
        copy: bool | None = ...,
    ) -> NativeDataFrameT | NativeSeriesT: ...

    def __call__(
        self,
        objs: Iterable[NativeDataFrameT | NativeSeriesT],
        *,
        axis: Axis,
        copy: bool | None = None,
    ) -> NativeDataFrameT | NativeSeriesT: ...
