"""Narwhals-level equivalent of `CompliantNamespace`."""

from __future__ import annotations

from typing import TYPE_CHECKING
from typing import Any
from typing import Callable
from typing import Generic
from typing import Literal
from typing import Protocol
from typing import TypeVar
from typing import overload

from narwhals._compliant.typing import CompliantNamespaceAny
from narwhals._compliant.typing import CompliantNamespaceT_co
from narwhals.dependencies import get_cudf
from narwhals.dependencies import get_modin
from narwhals.dependencies import get_pandas
from narwhals.dependencies import get_polars
from narwhals.dependencies import get_pyarrow
from narwhals.dependencies import is_dask_dataframe
from narwhals.dependencies import is_duckdb_relation
from narwhals.dependencies import is_pyspark_connect_dataframe
from narwhals.dependencies import is_pyspark_dataframe
from narwhals.dependencies import is_sqlframe_dataframe
from narwhals.utils import Implementation
from narwhals.utils import Version

if TYPE_CHECKING:
    from types import ModuleType
    from typing import ClassVar

    import duckdb
    import pandas as pd
    import polars as pl
    import pyarrow as pa
    import pyspark.sql as pyspark_sql
    from pyspark.sql.connect.dataframe import DataFrame as PySparkConnectDataFrame
    from typing_extensions import TypeAlias
    from typing_extensions import TypeIs

    from narwhals._arrow.namespace import ArrowNamespace
    from narwhals._dask.namespace import DaskNamespace
    from narwhals._duckdb.namespace import DuckDBNamespace
    from narwhals._pandas_like.namespace import PandasLikeNamespace
    from narwhals._polars.namespace import PolarsNamespace
    from narwhals._spark_like.dataframe import SQLFrameDataFrame
    from narwhals._spark_like.namespace import SparkLikeNamespace
    from narwhals.typing import DataFrameLike
    from narwhals.typing import NativeFrame
    from narwhals.typing import NativeLazyFrame
    from narwhals.typing import NativeSeries

    T = TypeVar("T")

    _Guard: TypeAlias = "Callable[[Any], TypeIs[T]]"

    _Polars: TypeAlias = Literal["polars"]
    _Arrow: TypeAlias = Literal["pyarrow"]
    _Dask: TypeAlias = Literal["dask"]
    _DuckDB: TypeAlias = Literal["duckdb"]
    _PandasLike: TypeAlias = Literal["pandas", "cudf", "modin"]
    _SparkLike: TypeAlias = Literal["pyspark", "sqlframe", "pyspark[connect]"]
    _EagerOnly: TypeAlias = "_PandasLike | _Arrow"
    _EagerAllowed: TypeAlias = "_Polars | _EagerOnly"
    _LazyOnly: TypeAlias = "_SparkLike | _Dask | _DuckDB"
    _LazyAllowed: TypeAlias = "_Polars | _LazyOnly"

    Polars: TypeAlias = Literal[_Polars, Implementation.POLARS]
    Arrow: TypeAlias = Literal[_Arrow, Implementation.PYARROW]
    Dask: TypeAlias = Literal[_Dask, Implementation.DASK]
    DuckDB: TypeAlias = Literal[_DuckDB, Implementation.DUCKDB]
    PandasLike: TypeAlias = Literal[
        _PandasLike, Implementation.PANDAS, Implementation.CUDF, Implementation.MODIN
    ]
    SparkLike: TypeAlias = Literal[
        _SparkLike,
        Implementation.PYSPARK,
        Implementation.SQLFRAME,
        Implementation.PYSPARK_CONNECT,
    ]
    EagerOnly: TypeAlias = "PandasLike | Arrow"
    EagerAllowed: TypeAlias = "EagerOnly | Polars"
    LazyOnly: TypeAlias = "SparkLike | Dask | DuckDB"
    LazyAllowed: TypeAlias = "LazyOnly | Polars"

    BackendName: TypeAlias = "_EagerAllowed | _LazyAllowed"
    IntoBackend: TypeAlias = "BackendName | Implementation | ModuleType"

    EagerAllowedNamespace: TypeAlias = "Namespace[PandasLikeNamespace] | Namespace[ArrowNamespace] | Namespace[PolarsNamespace]"
    EagerAllowedImplementation: TypeAlias = Literal[
        Implementation.PANDAS,
        Implementation.CUDF,
        Implementation.MODIN,
        Implementation.PYARROW,
        Implementation.POLARS,
    ]

    class _NativeDask(Protocol):
        _partition_type: type[pd.DataFrame]

    class _NativeCuDF(Protocol):
        def to_pylibcudf(self, *args: Any, **kwds: Any) -> Any: ...

    class _ModinDataFrame(Protocol):
        _pandas_class: type[pd.DataFrame]

    class _ModinSeries(Protocol):
        _pandas_class: type[pd.Series[Any]]

    _NativePolars: TypeAlias = "pl.DataFrame | pl.LazyFrame | pl.Series"
    _NativeArrow: TypeAlias = "pa.Table | pa.ChunkedArray[Any]"
    _NativeDuckDB: TypeAlias = "duckdb.DuckDBPyRelation"
    _NativePandas: TypeAlias = "pd.DataFrame | pd.Series[Any]"
    _NativeModin: TypeAlias = "_ModinDataFrame | _ModinSeries"
    _NativePandasLike: TypeAlias = "_NativePandas | _NativeCuDF | _NativeModin"
    _NativeSQLFrame: TypeAlias = "SQLFrameDataFrame"
    _NativePySpark: TypeAlias = "pyspark_sql.DataFrame"
    _NativePySparkConnect: TypeAlias = "PySparkConnectDataFrame"
    _NativeSparkLike: TypeAlias = (
        "_NativeSQLFrame | _NativePySpark | _NativePySparkConnect"
    )

    NativeKnown: TypeAlias = "_NativePolars | _NativeArrow | _NativePandasLike | _NativeSparkLike | _NativeDuckDB | _NativeDask"
    NativeUnknown: TypeAlias = (
        "NativeFrame | NativeSeries | NativeLazyFrame | DataFrameLike"
    )
    NativeAny: TypeAlias = "NativeKnown | NativeUnknown"

__all__ = ["Namespace"]


class Namespace(Generic[CompliantNamespaceT_co]):
    _compliant_namespace: CompliantNamespaceT_co
    _version: ClassVar[Version] = Version.MAIN

    def __init__(self, namespace: CompliantNamespaceT_co, /) -> None:
        self._compliant_namespace = namespace

    def __init_subclass__(cls, *args: Any, version: Version, **kwds: Any) -> None:
        super().__init_subclass__(*args, **kwds)

        if isinstance(version, Version):
            cls._version = version
        else:
            msg = f"Expected {Version} but got {type(version).__name__!r}"
            raise TypeError(msg)

    def __repr__(self) -> str:
        return f"Namespace[{type(self.compliant).__name__}]"

    @property
    def compliant(self) -> CompliantNamespaceT_co:
        return self._compliant_namespace

    @property
    def implementation(self) -> Implementation:
        return self.compliant._implementation

    @property
    def version(self) -> Version:
        return self._version

    @overload
    @classmethod
    def from_backend(cls, backend: PandasLike, /) -> Namespace[PandasLikeNamespace]: ...

    @overload
    @classmethod
    def from_backend(cls, backend: Polars, /) -> Namespace[PolarsNamespace]: ...

    @overload
    @classmethod
    def from_backend(cls, backend: Arrow, /) -> Namespace[ArrowNamespace]: ...

    @overload
    @classmethod
    def from_backend(cls, backend: SparkLike, /) -> Namespace[SparkLikeNamespace]: ...

    @overload
    @classmethod
    def from_backend(cls, backend: DuckDB, /) -> Namespace[DuckDBNamespace]: ...

    @overload
    @classmethod
    def from_backend(cls, backend: Dask, /) -> Namespace[DaskNamespace]: ...

    @overload
    @classmethod
    def from_backend(cls, backend: EagerAllowed, /) -> EagerAllowedNamespace: ...

    @overload
    @classmethod
    def from_backend(cls, backend: ModuleType, /) -> Namespace[CompliantNamespaceAny]: ...

    @classmethod
    def from_backend(
        cls: type[Namespace[Any]], backend: IntoBackend, /
    ) -> Namespace[Any]:
        """Instantiate from native namespace module, string, or Implementation.

        Arguments:
            backend: native namespace module, string, or Implementation.

        Returns:
            Namespace.

        Examples:
            >>> from narwhals._namespace import Namespace
            >>> Namespace.from_backend("polars")
            Namespace[PolarsNamespace]
        """
        impl = Implementation.from_backend(backend)
        backend_version = impl._backend_version()
        version = cls._version
        ns: CompliantNamespaceAny
        if impl.is_pandas_like():
            from narwhals._pandas_like.namespace import PandasLikeNamespace

            ns = PandasLikeNamespace(
                implementation=impl, backend_version=backend_version, version=version
            )

        elif impl.is_polars():
            from narwhals._polars.namespace import PolarsNamespace

            ns = PolarsNamespace(backend_version=backend_version, version=version)
        elif impl.is_pyarrow():
            from narwhals._arrow.namespace import ArrowNamespace

            ns = ArrowNamespace(backend_version=backend_version, version=version)
        elif impl.is_spark_like():
            from narwhals._spark_like.namespace import SparkLikeNamespace

            ns = SparkLikeNamespace(
                implementation=impl, backend_version=backend_version, version=version
            )
        elif impl.is_duckdb():
            from narwhals._duckdb.namespace import DuckDBNamespace

            ns = DuckDBNamespace(backend_version=backend_version, version=version)
        elif impl.is_dask():
            from narwhals._dask.namespace import DaskNamespace

            ns = DaskNamespace(backend_version=backend_version, version=version)
        else:
            msg = "Not supported Implementation"  # pragma: no cover
            raise AssertionError(msg)
        return cls(ns)

    @overload
    @classmethod
    def from_native_object(
        cls, native: _NativePolars, /
    ) -> Namespace[PolarsNamespace]: ...

    @overload
    @classmethod
    def from_native_object(
        cls, native: _NativePandasLike, /
    ) -> Namespace[PandasLikeNamespace]: ...

    @overload
    @classmethod
    def from_native_object(cls, native: _NativeArrow, /) -> Namespace[ArrowNamespace]: ...

    @overload
    @classmethod
    def from_native_object(
        cls, native: _NativeSparkLike, /
    ) -> Namespace[SparkLikeNamespace]: ...

    @overload
    @classmethod
    def from_native_object(
        cls, native: _NativeDuckDB, /
    ) -> Namespace[DuckDBNamespace]: ...

    @overload
    @classmethod
    def from_native_object(cls, native: _NativeDask, /) -> Namespace[DaskNamespace]: ...

    @overload
    @classmethod
    def from_native_object(
        cls, native: NativeUnknown, /
    ) -> Namespace[CompliantNamespaceAny]: ...

    @classmethod
    def from_native_object(  # noqa: PLR0911
        cls: type[Namespace[Any]], native: NativeAny, /
    ) -> Namespace[Any]:
        if is_native_polars(native):
            return cls.from_backend(Implementation.POLARS)
        elif is_native_pandas(native):
            return cls.from_backend(Implementation.PANDAS)
        elif is_native_arrow(native):
            return cls.from_backend(Implementation.PYARROW)
        elif is_native_spark_like(native):
            return cls.from_backend(
                Implementation.SQLFRAME
                if is_native_sqlframe(native)
                else Implementation.PYSPARK_CONNECT
                if is_native_pyspark_connect(native)
                else Implementation.PYSPARK
            )
        elif is_native_dask(native):
            return cls.from_backend(Implementation.DASK)  # pragma: no cover
        elif is_native_duckdb(native):
            return cls.from_backend(Implementation.DUCKDB)
        elif is_native_cudf(native):  # pragma: no cover
            return cls.from_backend(Implementation.CUDF)
        elif is_native_modin(native):  # pragma: no cover
            return cls.from_backend(Implementation.MODIN)
        else:
            msg = f"Unsupported type: {type(native).__qualname__!r}"
            raise TypeError(msg)


def is_native_polars(obj: Any) -> TypeIs[_NativePolars]:
    return (pl := get_polars()) is not None and isinstance(
        obj, (pl.DataFrame, pl.Series, pl.LazyFrame)
    )


def is_native_arrow(obj: Any) -> TypeIs[_NativeArrow]:
    return (pa := get_pyarrow()) is not None and isinstance(
        obj, (pa.Table, pa.ChunkedArray)
    )


def is_native_dask(obj: Any) -> TypeIs[_NativeDask]:
    return is_dask_dataframe(obj)


is_native_duckdb: _Guard[_NativeDuckDB] = is_duckdb_relation
is_native_sqlframe: _Guard[_NativeSQLFrame] = is_sqlframe_dataframe
is_native_pyspark: _Guard[_NativePySpark] = is_pyspark_dataframe
is_native_pyspark_connect: _Guard[_NativePySparkConnect] = is_pyspark_connect_dataframe


def is_native_pandas(obj: Any) -> TypeIs[_NativePandas]:
    return (pd := get_pandas()) is not None and isinstance(obj, (pd.DataFrame, pd.Series))


def is_native_modin(obj: Any) -> TypeIs[_NativeModin]:
    return (mpd := get_modin()) is not None and isinstance(
        obj, (mpd.DataFrame, mpd.Series)
    )  # pragma: no cover


def is_native_cudf(obj: Any) -> TypeIs[_NativeCuDF]:
    return (cudf := get_cudf()) is not None and isinstance(
        obj, (cudf.DataFrame, cudf.Series)
    )  # pragma: no cover


def is_native_pandas_like(obj: Any) -> TypeIs[_NativePandasLike]:
    return (
        is_native_pandas(obj) or is_native_cudf(obj) or is_native_modin(obj)
    )  # pragma: no cover


def is_native_spark_like(obj: Any) -> TypeIs[_NativeSparkLike]:
    return (
        is_native_sqlframe(obj)
        or is_native_pyspark(obj)
        or is_native_pyspark_connect(obj)
    )