"""
Backend for ``alt.datasets.Loader``.

Notes
-----
Extending would be more ergonomic if `read`, `scan`, `_constraints` were available under a single export::

    from altair.datasets import ext, reader
    import polars as pl

    impls = (
        ext.read(pl.read_parquet, ext.is_parquet),
        ext.read(pl.read_csv, ext.is_csv),
        ext.read(pl.read_json, ext.is_json),
    )
    user_reader = reader(impls)
    user_reader.dataset("airports")
"""

from __future__ import annotations

from collections import Counter
from collections.abc import Mapping
from importlib import import_module
from importlib.util import find_spec
from itertools import chain
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, overload
from urllib.request import build_opener as _build_opener

from narwhals.stable import v1 as nw
from packaging.requirements import Requirement

from altair.datasets import _readimpl
from altair.datasets._cache import CsvCache, DatasetCache, SchemaCache, _iter_metadata
from altair.datasets._constraints import is_parquet
from altair.datasets._exceptions import AltairDatasetsError, module_not_found
from altair.datasets._readimpl import IntoDataFrameT, IntoLazyFrameT, is_available

if TYPE_CHECKING:
    import sys
    from collections.abc import Callable, Sequence
    from typing import LiteralString, TypeAlias
    from urllib.request import OpenerDirector

    import pandas as pd
    import polars as pl
    import pyarrow as pa
    from narwhals.stable.v1.typing import IntoExpr

    from altair.datasets._readimpl import BaseImpl, R, Read, Scan
    from altair.datasets._typing import Dataset, Extension, Metadata
    from altair.vegalite.v6.schema._typing import OneOrSeq

    if sys.version_info >= (3, 13):
        from typing import TypeIs, TypeVar
    else:
        from typing_extensions import TypeIs, TypeVar
    if sys.version_info >= (3, 12):
        from typing import Unpack
    else:
        from typing_extensions import Unpack
    _Polars: TypeAlias = Literal["polars"]
    _Pandas: TypeAlias = Literal["pandas"]
    _PyArrow: TypeAlias = Literal["pyarrow"]
    _PandasAny: TypeAlias = Literal[_Pandas, "pandas[pyarrow]"]
    _Backend: TypeAlias = Literal[_Polars, _PandasAny, _PyArrow]
    _CuDF: TypeAlias = Literal["cudf"]
    _Dask: TypeAlias = Literal["dask"]
    _DuckDB: TypeAlias = Literal["duckdb"]
    _Ibis: TypeAlias = Literal["ibis"]
    _PySpark: TypeAlias = Literal["pyspark"]
    _NwSupport: TypeAlias = Literal[
        _Polars, _Pandas, _PyArrow, _CuDF, _Dask, _DuckDB, _Ibis, _PySpark
    ]
    _NwSupportT = TypeVar(
        "_NwSupportT",
        _Polars,
        _Pandas,
        _PyArrow,
        _CuDF,
        _Dask,
        _DuckDB,
        _Ibis,
        _PySpark,
    )
    _EagerAllowedImpl: TypeAlias = Literal[
        nw.Implementation.PANDAS,
        nw.Implementation.POLARS,
        nw.Implementation.PYARROW,
    ]
    _EagerAllowed: TypeAlias = Literal[_Pandas, _Polars, _PyArrow]

_SupportProfile: TypeAlias = Mapping[
    Literal["supported", "unsupported"], "Sequence[Dataset]"
]
"""
Dataset support varies between backends and available dependencies.

Any name listed in ``"unsupported"`` will raise an error on::

    from altair.datasets import load

    load("7zip")

Instead, they can be loaded via::

    import altair as alt
    from altair.datasets import url

    alt.Chart(url("7zip"))
"""


class Reader(Generic[IntoDataFrameT, IntoLazyFrameT]):
    """
    Modular file reader, targeting remote & local tabular resources.

    .. warning::
        Use ``reader(...)`` instead of instantiating ``Reader`` directly.
    """

    _read: Sequence[Read[IntoDataFrameT]]
    """Eager file read functions."""

    _scan: Sequence[Scan[IntoLazyFrameT]]
    """Lazy file read functions."""

    _name: str
    """
    Used in error messages, repr and matching ``@overload``(s).

    Otherwise, has no concrete meaning.
    """

    _implementation: _EagerAllowedImpl
    """
    Corresponding `narwhals implementation`_.

    .. _narwhals implementation:
        https://github.com/narwhals-dev/narwhals/blob/9b6a355530ea46c590d5a6d1d0567be59c0b5742/narwhals/utils.py#L61-L290
    """

    _opener: ClassVar[OpenerDirector] = _build_opener()
    _metadata_path: ClassVar[Path] = (
        Path(__file__).parent / "_metadata" / "metadata.parquet"
    )

    def __init__(
        self,
        read: Sequence[Read[IntoDataFrameT]],
        scan: Sequence[Scan[IntoLazyFrameT]],
        name: str,
        implementation: _EagerAllowedImpl,
    ) -> None:
        self._read = read
        self._scan = scan
        self._name = name
        self._implementation = implementation
        self._schema_cache = SchemaCache(implementation=implementation)

    def __repr__(self) -> str:
        from textwrap import indent

        PREFIX = " " * 4
        NL = "\n"
        body = f"read\n{indent(NL.join(str(el) for el in self._read), PREFIX)}"
        if self._scan:
            body += f"\nscan\n{indent(NL.join(str(el) for el in self._scan), PREFIX)}"
        return f"Reader[{self._name}] {self._implementation!r}\n{body}"

    def read_fn(self, meta: Metadata, /) -> Callable[..., IntoDataFrameT]:
        return self._solve(meta, self._read)

    def scan_fn(self, meta: Metadata | Path | str, /) -> Callable[..., IntoLazyFrameT]:
        meta = meta if isinstance(meta, Mapping) else {"suffix": _into_suffix(meta)}
        return self._solve(meta, self._scan)

    @property
    def cache(self) -> DatasetCache:
        return DatasetCache(self)

    def _handle_pyarrow_date_error(self, e: Exception, name: str) -> None:
        """Handle PyArrow date parsing errors with informative error messages, see https://github.com/apache/arrow/issues/41488."""
        if "CSV conversion error to date" in str(e) and "pyarrow" in str(
            type(e).__module__
        ):
            message = (
                f"PyArrow cannot parse date format in dataset '{name}'. "
                f"This is a known limitation of PyArrow's CSV reader for non-ISO date formats.\n\n"
                f"Alternatives:\n"
                f"1. Use a different backend: data.{name}(engine='pandas') or data.{name}(engine='polars')\n"
                f"2. Convert dates manually after loading as strings\n\n"
                f"Original error: {e}"
            )
            raise AltairDatasetsError(message) from e
        raise e

    def dataset(
        self,
        name: Dataset | LiteralString,
        suffix: Extension | None = None,
        /,
        **kwds: Any,
    ) -> IntoDataFrameT:
        frame = self._query(name, suffix)
        meta = next(_iter_metadata(frame))
        fn = self.read_fn(meta)
        fn_kwds = self._merge_kwds(meta, kwds)
        if self.cache.is_active():
            fp = self.cache._maybe_download(meta)
            try:
                return fn(fp, **fn_kwds)
            except Exception as e:
                self._handle_pyarrow_date_error(e, name)
                raise
        else:
            with self._opener.open(meta["url"]) as f:
                try:
                    return fn(f, **fn_kwds)
                except Exception as e:
                    self._handle_pyarrow_date_error(e, name)
                    raise

    def url(
        self, name: Dataset | LiteralString, suffix: Extension | None = None, /
    ) -> str:
        frame = self._query(name, suffix)
        meta = next(_iter_metadata(frame))
        if is_parquet(meta.items()) and not is_available("vegafusion"):
            raise AltairDatasetsError.from_url(meta)
        url = meta["url"]
        if isinstance(url, str):
            return url
        else:
            msg = f"Expected 'str' but got {type(url).__name__!r}\nfrom {url!r}."
            raise TypeError(msg)

    # TODO: (Multiple)
    # - Settle on a better name
    # - Add method to `Loader`
    # - Move docs to `Loader.{new name}`
    def open_markdown(self, name: Dataset, /) -> None:
        """
        Learn more about a dataset, opening `vega-datasets/datapackage.md`_ with the default browser.

        Additional info *may* include: `description`_, `schema`_, `sources`_, `licenses`_.

        .. _vega-datasets/datapackage.md:
            https://github.com/vega/vega-datasets/blob/main/datapackage.md
        .. _description:
            https://datapackage.org/standard/data-resource/#description
        .. _schema:
            https://datapackage.org/standard/table-schema/#schema
        .. _sources:
            https://datapackage.org/standard/data-package/#sources
        .. _licenses:
            https://datapackage.org/standard/data-package/#licenses
        """
        import webbrowser

        from altair.utils import VERSIONS

        ref = self._query(name).get_column("file_name").item(0).replace(".", "")
        tag = VERSIONS["vega-datasets"]
        url = f"https://github.com/vega/vega-datasets/blob/v{tag}/datapackage.md#{ref}"
        webbrowser.open(url)

    @overload
    def profile(self, *, show: Literal[False] = ...) -> _SupportProfile: ...

    @overload
    def profile(self, *, show: Literal[True]) -> None: ...

    def profile(self, *, show: bool = False) -> _SupportProfile | None:
        """
        Describe which datasets can be loaded as tabular data.

        Parameters
        ----------
        show
            Print a densely formatted repr *instead of* returning a mapping.
        """
        relevant_columns = set(
            chain.from_iterable(impl._relevant_columns for impl in self._read)
        )
        frame = self._scan_metadata().select("dataset_name", *relevant_columns)
        inc_expr = nw.any_horizontal(impl._include_expr for impl in self._read)
        result: _SupportProfile = {
            "unsupported": _dataset_names(frame, ~inc_expr),
            "supported": _dataset_names(frame, inc_expr),
        }
        if show:
            import pprint

            pprint.pprint(result, compact=True, sort_dicts=False)
            return None
        return result

    def _query(
        self, name: Dataset | LiteralString, suffix: Extension | None = None, /
    ) -> nw.DataFrame[IntoDataFrameT]:
        """
        Query a tabular version of `vega-datasets/datapackage.json`_.

        Applies a filter, erroring out when no results would be returned.

        .. _vega-datasets/datapackage.json:
            https://github.com/vega/vega-datasets/blob/main/datapackage.json
        """
        constraints = _into_constraints(name, suffix)
        frame = self._scan_metadata(**constraints).collect()
        if not frame.is_empty():
            return frame
        else:
            msg = f"Found no results for:\n    {constraints!r}"
            raise ValueError(msg)

    def _merge_kwds(self, meta: Metadata, kwds: dict[str, Any], /) -> Mapping[str, Any]:
        """
        Extend user-provided arguments with dataset & library-specfic defaults.

        .. important:: User-provided arguments have a higher precedence.
        """
        if self._schema_cache.is_active() and (
            schema := self._schema_cache.schema_kwds(meta)
        ):
            kwds = schema | kwds if kwds else schema
        return kwds

    @property
    def _metadata_frame(self) -> nw.LazyFrame[IntoLazyFrameT]:
        fp = self._metadata_path
        return nw.from_native(self.scan_fn(fp)(fp)).lazy()

    def _scan_metadata(
        self, *predicates: OneOrSeq[IntoExpr], **constraints: Unpack[Metadata]
    ) -> nw.LazyFrame[IntoLazyFrameT]:
        if predicates or constraints:
            return self._metadata_frame.filter(*predicates, **constraints)
        return self._metadata_frame

    def _solve(
        self, meta: Metadata, impls: Sequence[BaseImpl[R]], /
    ) -> Callable[..., R]:
        """
        Return the first function that satisfies dataset constraints.

        See Also
        --------
        ``altair.datasets._readimpl.BaseImpl.unwrap_or_skip``
        """
        items = meta.items()
        it = (some for impl in impls if (some := impl.unwrap_or_skip(items)))
        if fn_or_err := next(it, None):
            if _is_err(fn_or_err):
                raise fn_or_err.from_tabular(meta, self._name)
            return fn_or_err
        raise AltairDatasetsError.from_tabular(meta, self._name)


def _dataset_names(
    frame: nw.LazyFrame, *predicates: OneOrSeq[IntoExpr]
) -> Sequence[Dataset]:
    # NOTE: helper function for `Reader.profile`
    return (
        frame.filter(*predicates)
        .select("dataset_name")
        .collect()
        .get_column("dataset_name")
        .to_list()
    )


class _NoParquetReader(Reader[IntoDataFrameT]):
    def __repr__(self) -> str:
        return f"{super().__repr__()}\ncsv_cache\n    {self.csv_cache!r}"

    @property
    def csv_cache(self) -> CsvCache:
        if not hasattr(self, "_csv_cache"):
            self._csv_cache = CsvCache()
        return self._csv_cache

    @property
    def _metadata_frame(self) -> nw.LazyFrame[Any]:
        data = self.csv_cache.rotated
        impl = self._implementation
        return nw.maybe_convert_dtypes(nw.from_dict(data, backend=impl)).lazy()


@overload
def reader(
    read_fns: Sequence[Read[IntoDataFrameT]],
    scan_fns: tuple[()] = ...,
    *,
    name: str | None = ...,
    implementation: nw.Implementation = ...,
) -> Reader[IntoDataFrameT]: ...


@overload
def reader(
    read_fns: Sequence[Read[IntoDataFrameT]],
    scan_fns: Sequence[Scan[IntoLazyFrameT]],
    *,
    name: str | None = ...,
    implementation: nw.Implementation = ...,
) -> Reader[IntoDataFrameT, IntoLazyFrameT]: ...


def reader(
    read_fns: Sequence[Read[IntoDataFrameT]],
    scan_fns: Sequence[Scan[IntoLazyFrameT]] = (),
    *,
    name: str | None = None,
    implementation: nw.Implementation = nw.Implementation.UNKNOWN,
) -> Reader[IntoDataFrameT, IntoLazyFrameT] | Reader[IntoDataFrameT]:
    name = name or Counter(el._inferred_package for el in read_fns).most_common(1)[0][0]
    if not _is_eager_allowed(implementation):
        implementation = _into_implementation(Requirement(name))
    if scan_fns:
        return Reader(read_fns, scan_fns, name, implementation)
    if stolen := _steal_eager_parquet(read_fns):
        return Reader(read_fns, stolen, name, implementation)
    else:
        return _NoParquetReader[IntoDataFrameT](read_fns, (), name, implementation)


def infer_backend(
    *, priority: Sequence[_Backend] = ("polars", "pandas[pyarrow]", "pandas", "pyarrow")
) -> Reader[Any, Any]:
    """
    Return the first available reader in order of `priority`.

    Notes
    -----
    - ``"polars"``: can natively load every dataset (including ``(Geo|Topo)JSON``)
    - ``"pandas[pyarrow]"``: can load *most* datasets, guarantees ``.parquet`` support
    - ``"pandas"``: supports ``.parquet``, if `fastparquet`_ is installed
    - ``"pyarrow"``: least reliable

    .. _fastparquet:
        https://github.com/dask/fastparquet
    """
    it = (_from_backend(name) for name in priority if is_available(_requirements(name)))
    if reader := next(it, None):
        return reader
    raise AltairDatasetsError.from_priority(priority)


@overload
def _from_backend(name: _Polars, /) -> Reader[pl.DataFrame, pl.LazyFrame]: ...
@overload
def _from_backend(name: _PandasAny, /) -> Reader[pd.DataFrame]: ...
@overload
def _from_backend(name: _PyArrow, /) -> Reader[pa.Table]: ...


# FIXME: The order this is defined in makes splitting the module complicated
# - Can't use a classmethod, since some result in a subclass used
def _from_backend(name: _Backend, /) -> Reader[Any, Any]:
    """
    Reader initialization dispatcher.

    FIXME: Works, but defining these in mixed shape functions seems off.
    """
    if not _is_backend(name):
        msg = f"Unknown backend {name!r}"
        raise TypeError(msg)
    implementation = _into_implementation(name)
    if name == "polars":
        rd, sc = _readimpl.pl_only()
        return reader(rd, sc, name=name, implementation=implementation)
    elif name == "pandas[pyarrow]":
        return reader(_readimpl.pd_pyarrow(), name=name, implementation=implementation)
    elif name == "pandas":
        return reader(_readimpl.pd_only(), name=name, implementation=implementation)
    elif name == "pyarrow":
        return reader(_readimpl.pa_any(), name=name, implementation=implementation)


def _is_backend(obj: Any) -> TypeIs[_Backend]:
    return obj in {"polars", "pandas", "pandas[pyarrow]", "pyarrow"}


def _is_err(obj: Any) -> TypeIs[type[AltairDatasetsError]]:
    return obj is AltairDatasetsError


def _into_constraints(
    name: Dataset | LiteralString, suffix: Extension | None, /
) -> Metadata:
    """Transform args into a mapping to column names."""
    m: Metadata = {}
    if "." in name:
        m["file_name"] = name
    elif suffix is None:
        m["dataset_name"] = name
    elif suffix.startswith("."):
        m = {"dataset_name": name, "suffix": suffix}
    else:
        from typing import get_args

        from altair.datasets._typing import Extension

        msg = (
            f"Expected 'suffix' to be one of {get_args(Extension)!r},\n"
            f"but got: {suffix!r}"
        )
        raise TypeError(msg)
    return m


def _is_eager_allowed(impl: nw.Implementation, /) -> TypeIs[_EagerAllowedImpl]:
    return impl in {
        nw.Implementation.PANDAS,
        nw.Implementation.POLARS,
        nw.Implementation.PYARROW,
    }


def _into_implementation(
    backend: _NwSupport | _PandasAny | nw.Implementation | Requirement, /
) -> _EagerAllowedImpl:
    req = (
        Requirement(str(backend)) if isinstance(backend, nw.Implementation) else backend
    )
    primary = _import_guarded(req)
    impl = nw.Implementation.from_string(primary)
    if not _is_eager_allowed(impl):
        if impl is nw.Implementation.UNKNOWN:
            msg = f"Package {primary!r} is not supported by `narwhals`."
            raise ValueError(msg)
        raise NotImplementedError(impl)
    return impl


def _into_suffix(obj: Path | str, /) -> Any:
    if isinstance(obj, Path):
        return obj.suffix
    elif isinstance(obj, str):
        return obj
    else:
        msg = f"Unexpected type {type(obj).__name__!r}"
        raise TypeError(msg)


def _steal_eager_parquet(
    read_fns: Sequence[Read[IntoDataFrameT]], /
) -> Sequence[Scan[Any]] | None:
    if convertible := next((rd for rd in read_fns if rd.include <= is_parquet), None):
        return (_readimpl.into_scan(convertible),)
    return None


@overload
def _import_guarded(req: _PandasAny, /) -> _Pandas: ...


@overload
def _import_guarded(req: _NwSupportT, /) -> _NwSupportT: ...


@overload
def _import_guarded(req: Requirement, /) -> LiteralString: ...


def _import_guarded(req: Any, /) -> LiteralString:
    requires = _requirements(req)
    for name in requires:
        if spec := find_spec(name):
            import_module(spec.name)
        else:
            raise module_not_found(str(req), requires, missing=name)
    return requires[0]


def _requirements(req: Requirement | str, /) -> tuple[Any, ...]:
    req = Requirement(req) if isinstance(req, str) else req
    return (req.name, *req.extras)