from __future__ import annotations

import operator
from functools import lru_cache, partial
from typing import TYPE_CHECKING, Any, Literal, cast, overload

import ibis
import ibis.expr.datatypes as ibis_dtypes

from narwhals._utils import Version, isinstance_or_issubclass

if TYPE_CHECKING:
    from collections.abc import Callable, Mapping
    from typing import TypeAlias

    import ibis.expr.types as ir
    from ibis.common.temporal import TimestampUnit
    from ibis.expr.datatypes import DataType as IbisDataType
    from typing_extensions import TypeIs

    from narwhals._duration import IntervalUnit
    from narwhals._ibis.dataframe import IbisLazyFrame
    from narwhals._ibis.expr import IbisExpr
    from narwhals.dtypes import DType
    from narwhals.typing import IntoDType, PythonLiteral

IntoColumn: TypeAlias = "str | ir.Value | ir.Column"
SortFn: TypeAlias = "Callable[[IntoColumn], ir.Column]"
Incomplete: TypeAlias = Any
"""Marker for upstream issues."""


@overload
def lit(value: bool, dtype: None = ...) -> ir.BooleanScalar: ...  # noqa: FBT001
@overload
def lit(value: int, dtype: None = ...) -> ir.IntegerScalar: ...
@overload
def lit(value: float, dtype: None = ...) -> ir.FloatingScalar: ...
@overload
def lit(value: str, dtype: None = ...) -> ir.StringScalar: ...
@overload
def lit(value: PythonLiteral | ir.Value, dtype: None = ...) -> ir.Scalar: ...
@overload
def lit(value: Any, dtype: Any) -> Incomplete: ...
def lit(value: Any, dtype: Any | None = None) -> Incomplete:
    """Alias for `ibis.literal`."""
    literal: Incomplete = ibis.literal
    return literal(value, dtype)


asc_nulls_first = cast("SortFn", partial(ibis.asc, nulls_first=True))
asc_nulls_last = cast("SortFn", partial(ibis.asc, nulls_first=False))
desc_nulls_first = cast("SortFn", partial(ibis.desc, nulls_first=True))
desc_nulls_last = cast("SortFn", partial(ibis.desc, nulls_first=False))


BucketUnit: TypeAlias = Literal[
    "years",
    "quarters",
    "months",
    "days",
    "hours",
    "minutes",
    "seconds",
    "milliseconds",
    "microseconds",
    "nanoseconds",
]
TruncateUnit: TypeAlias = Literal[
    "Y", "Q", "M", "W", "D", "h", "m", "s", "ms", "us", "ns"
]

UNITS_DICT_BUCKET: Mapping[IntervalUnit, BucketUnit] = {
    "y": "years",
    "q": "quarters",
    "mo": "months",
    "d": "days",
    "h": "hours",
    "m": "minutes",
    "s": "seconds",
    "ms": "milliseconds",
    "us": "microseconds",
    "ns": "nanoseconds",
}

UNITS_DICT_TRUNCATE: Mapping[IntervalUnit, TruncateUnit] = {
    "y": "Y",
    "q": "Q",
    "mo": "M",
    "d": "D",
    "h": "h",
    "m": "m",
    "s": "s",
    "ms": "ms",
    "us": "us",
    "ns": "ns",
}

FUNCTION_REMAPPING = {
    "starts_with": "startswith",
    "ends_with": "endswith",
    "regexp_matches": "re_search",
    "str_split": "split",
    "dayofyear": "day_of_year",
    "to_date": "date",
    "bool_and": "all",
    "bool_or": "any",
    "isnotnull": "notnull",
}


def evaluate_exprs(df: IbisLazyFrame, /, *exprs: IbisExpr) -> list[tuple[str, ir.Value]]:
    native_results: list[tuple[str, ir.Value]] = []
    for expr in exprs:
        native_series_list = expr(df)
        output_names = expr._evaluate_output_names(df)
        if expr._alias_output_names is not None:
            output_names = expr._alias_output_names(output_names)
        if len(output_names) != len(native_series_list):  # pragma: no cover
            msg = f"Internal error: got output names {output_names}, but only got {len(native_series_list)} results"
            raise AssertionError(msg)
        native_results.extend(zip(output_names, native_series_list, strict=False))
    return native_results


@lru_cache(maxsize=16)
def native_to_narwhals_dtype(ibis_dtype: IbisDataType, version: Version) -> DType:  # noqa: C901, PLR0912
    dtypes = version.dtypes
    if ibis_dtype.is_int64():
        return dtypes.Int64()
    if ibis_dtype.is_int32():
        return dtypes.Int32()
    if ibis_dtype.is_int16():
        return dtypes.Int16()
    if ibis_dtype.is_int8():
        return dtypes.Int8()
    if ibis_dtype.is_uint64():
        return dtypes.UInt64()
    if ibis_dtype.is_uint32():
        return dtypes.UInt32()
    if ibis_dtype.is_uint16():
        return dtypes.UInt16()
    if ibis_dtype.is_uint8():
        return dtypes.UInt8()
    if ibis_dtype.is_boolean():
        return dtypes.Boolean()
    if ibis_dtype.is_float64():
        return dtypes.Float64()
    if ibis_dtype.is_float32():
        return dtypes.Float32()
    if ibis_dtype.is_float16():
        return dtypes.Float16()
    if ibis_dtype.is_string():
        return dtypes.String()
    if ibis_dtype.is_date():
        return dtypes.Date()
    if is_timestamp(ibis_dtype):
        _unit = cast("TimestampUnit", ibis_dtype.unit)
        return dtypes.Datetime(time_unit=_unit.value, time_zone=ibis_dtype.timezone)
    if is_interval(ibis_dtype):
        _time_unit = ibis_dtype.unit.value
        if _time_unit not in {"ns", "us", "ms", "s"}:  # pragma: no cover
            msg = f"Unsupported interval unit: {_time_unit}"
            raise NotImplementedError(msg)
        return dtypes.Duration(_time_unit)
    if is_array(ibis_dtype):
        if ibis_dtype.length:
            return dtypes.Array(
                native_to_narwhals_dtype(ibis_dtype.value_type, version),
                ibis_dtype.length,
            )
        return dtypes.List(native_to_narwhals_dtype(ibis_dtype.value_type, version))
    if is_struct(ibis_dtype):
        return dtypes.Struct(
            [
                dtypes.Field(name, native_to_narwhals_dtype(dtype, version))
                for name, dtype in ibis_dtype.items()
            ]
        )
    if is_decimal(ibis_dtype):
        # Same default as in ibis.Decimal.{to_polars, to_pyarrow}
        # https://github.com/ibis-project/ibis/blob/5028d8e5e1921d5de48f59fad48b86c0de541b0d/ibis/formats/polars.py#L82-L86
        # https://github.com/ibis-project/ibis/blob/5028d8e5e1921d5de48f59fad48b86c0de541b0d/ibis/formats/pyarrow.py#L165-L177
        # According to their own comment:
        # > set default precision and scale to something; unclear how to choose this
        # source: (see https://github.com/ibis-project/ibis/blob/5028d8e5e1921d5de48f59fad48b86c0de541b0d/ibis/formats/pyarrow.py#L166)
        scale = 9 if ibis_dtype.scale is None else ibis_dtype.scale
        return dtypes.Decimal(precision=ibis_dtype.precision, scale=scale)
    if ibis_dtype.is_time():
        return dtypes.Time()
    if ibis_dtype.is_binary():
        return dtypes.Binary()
    return dtypes.Unknown()  # pragma: no cover


def is_timestamp(obj: IbisDataType) -> TypeIs[ibis_dtypes.Timestamp]:
    return obj.is_timestamp()


def is_interval(obj: IbisDataType) -> TypeIs[ibis_dtypes.Interval]:
    return obj.is_interval()


def is_array(obj: IbisDataType) -> TypeIs[ibis_dtypes.Array[Any]]:
    return obj.is_array()


def is_struct(obj: IbisDataType) -> TypeIs[ibis_dtypes.Struct]:
    return obj.is_struct()


def is_floating(obj: IbisDataType) -> TypeIs[ibis_dtypes.Floating]:
    return obj.is_floating()


def is_decimal(obj: IbisDataType) -> TypeIs[ibis_dtypes.Decimal]:
    return obj.is_decimal()


dtypes = Version.MAIN.dtypes
NW_TO_IBIS_DTYPES: Mapping[type[DType], IbisDataType] = {
    dtypes.Float64: ibis_dtypes.Float64(),
    dtypes.Float32: ibis_dtypes.Float32(),
    dtypes.Float16: ibis_dtypes.Float16(),
    dtypes.Binary: ibis_dtypes.Binary(),
    dtypes.String: ibis_dtypes.String(),
    dtypes.Boolean: ibis_dtypes.Boolean(),
    dtypes.Date: ibis_dtypes.Date(),
    dtypes.Time: ibis_dtypes.Time(),
    dtypes.Int8: ibis_dtypes.Int8(),
    dtypes.Int16: ibis_dtypes.Int16(),
    dtypes.Int32: ibis_dtypes.Int32(),
    dtypes.Int64: ibis_dtypes.Int64(),
    dtypes.UInt8: ibis_dtypes.UInt8(),
    dtypes.UInt16: ibis_dtypes.UInt16(),
    dtypes.UInt32: ibis_dtypes.UInt32(),
    dtypes.UInt64: ibis_dtypes.UInt64(),
}
# Enum support: https://github.com/ibis-project/ibis/issues/10991
UNSUPPORTED_DTYPES = (dtypes.Int128, dtypes.UInt128, dtypes.Categorical, dtypes.Enum)


def narwhals_to_native_dtype(dtype: IntoDType, version: Version) -> IbisDataType:
    dtypes = version.dtypes
    base_type = dtype.base_type()
    if ibis_type := NW_TO_IBIS_DTYPES.get(base_type):
        return ibis_type
    if isinstance_or_issubclass(dtype, dtypes.Datetime):
        return ibis_dtypes.Timestamp.from_unit(dtype.time_unit, timezone=dtype.time_zone)
    if isinstance_or_issubclass(dtype, dtypes.Duration):
        return ibis_dtypes.Interval(unit=dtype.time_unit)  # pyright: ignore[reportArgumentType]
    if isinstance_or_issubclass(dtype, dtypes.List):
        inner = narwhals_to_native_dtype(dtype.inner, version)
        return ibis_dtypes.Array(value_type=inner)
    if isinstance_or_issubclass(dtype, dtypes.Struct):
        fields = [
            (field.name, narwhals_to_native_dtype(field.dtype, version))
            for field in dtype.fields
        ]
        return ibis_dtypes.Struct.from_tuples(fields)
    if isinstance_or_issubclass(dtype, dtypes.Array):
        inner = narwhals_to_native_dtype(dtype.inner, version)
        return ibis_dtypes.Array(value_type=inner, length=dtype.size)
    if isinstance_or_issubclass(dtype, dtypes.Decimal):
        return ibis_dtypes.Decimal(dtype.precision, dtype.scale)
    if issubclass(base_type, UNSUPPORTED_DTYPES):
        msg = f"Converting to {base_type.__name__} dtype is not supported for Ibis."
        raise NotImplementedError(msg)
    msg = f"Unknown dtype: {dtype}"  # pragma: no cover
    raise AssertionError(msg)


_BINARY_OPS = {
    "add": operator.add,
    "subtract": operator.sub,
    "multiply": operator.mul,
    "divide": operator.truediv,
    "floordiv": operator.floordiv,
    "and": operator.and_,
    "or": operator.or_,
}


def function(name: str, *args: ir.Value | PythonLiteral) -> ir.Value:
    # Workaround SQL vs Ibis differences.
    if name in _BINARY_OPS:
        return _BINARY_OPS[name](*args)
    if name == "row_number":
        return ibis.row_number() + lit(1)
    if name == "least":
        return ibis.least(*args)
    if name == "greatest":
        return ibis.greatest(*args)
    expr = args[0]
    if name == "var_pop":
        return cast("ir.NumericColumn", expr).var(how="pop")
    if name == "var_samp":
        return cast("ir.NumericColumn", expr).var(how="sample")
    if name == "stddev_pop":
        return cast("ir.NumericColumn", expr).std(how="pop")
    if name == "stddev_samp":
        return cast("ir.NumericColumn", expr).std(how="sample")
    if name == "substr":
        # Ibis is 0-indexed here, SQL is 1-indexed
        return cast("ir.StringColumn", expr).substr(args[1] - 1, *args[2:])  # type: ignore[operator]  # pyright: ignore[reportArgumentType]
    return getattr(expr, FUNCTION_REMAPPING.get(name, name))(*args[1:])
