"""
Column — wraps a polars.Expr, mirrors pyspark.sql.Column API.
"""
from __future__ import annotations

import polars as pl


def _to_expr(val) -> pl.Expr:
    """Convert a Column, str column name, or literal value to a pl.Expr."""
    if isinstance(val, Column):
        return val._expr
    if isinstance(val, pl.Expr):
        return val
    return pl.lit(val)


def _to_col(val) -> "Column":
    if isinstance(val, Column):
        return val
    if isinstance(val, str):
        return Column(pl.col(val))
    return Column(pl.lit(val))


class Column:
    """
    Mirrors pyspark.sql.Column. Wraps a polars.Expr so that PySpark-style
    column operations translate directly to Polars lazy expressions.
    """

    def __init__(self, expr: pl.Expr | str):
        if isinstance(expr, str):
            self._expr: pl.Expr = pl.col(expr)
        else:
            self._expr = expr

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _wrap(self, expr: pl.Expr) -> "Column":
        return Column(expr)

    # ------------------------------------------------------------------
    # Arithmetic
    # ------------------------------------------------------------------

    def __add__(self, other):  return self._wrap(self._expr + _to_expr(other))
    def __radd__(self, other): return self._wrap(_to_expr(other) + self._expr)
    def __sub__(self, other):  return self._wrap(self._expr - _to_expr(other))
    def __rsub__(self, other): return self._wrap(_to_expr(other) - self._expr)
    def __mul__(self, other):  return self._wrap(self._expr * _to_expr(other))
    def __rmul__(self, other): return self._wrap(_to_expr(other) * self._expr)
    def __truediv__(self, other):  return self._wrap(self._expr / _to_expr(other))
    def __rtruediv__(self, other): return self._wrap(_to_expr(other) / self._expr)
    def __floordiv__(self, other): return self._wrap(self._expr // _to_expr(other))
    def __mod__(self, other):  return self._wrap(self._expr % _to_expr(other))
    def __pow__(self, other):  return self._wrap(self._expr ** _to_expr(other))
    def __neg__(self):         return self._wrap(-self._expr)

    # ------------------------------------------------------------------
    # Comparison
    # ------------------------------------------------------------------

    def __eq__(self, other):  return self._wrap(self._expr == _to_expr(other))  # type: ignore[override]
    def __ne__(self, other):  return self._wrap(self._expr != _to_expr(other))  # type: ignore[override]
    def __lt__(self, other):  return self._wrap(self._expr < _to_expr(other))
    def __le__(self, other):  return self._wrap(self._expr <= _to_expr(other))
    def __gt__(self, other):  return self._wrap(self._expr > _to_expr(other))
    def __ge__(self, other):  return self._wrap(self._expr >= _to_expr(other))

    def __hash__(self):  # needed because we override __eq__
        return id(self)

    # ------------------------------------------------------------------
    # Logical
    # ------------------------------------------------------------------

    def __and__(self, other):  return self._wrap(self._expr & _to_expr(other))
    def __rand__(self, other): return self._wrap(_to_expr(other) & self._expr)
    def __or__(self, other):   return self._wrap(self._expr | _to_expr(other))
    def __ror__(self, other):  return self._wrap(_to_expr(other) | self._expr)
    def __invert__(self):      return self._wrap(~self._expr)

    # ------------------------------------------------------------------
    # Null checks
    # ------------------------------------------------------------------

    def isNull(self) -> "Column":    return self._wrap(self._expr.is_null())
    def isNotNull(self) -> "Column": return self._wrap(self._expr.is_not_null())
    def isNaN(self) -> "Column":     return self._wrap(self._expr.is_nan())
    def isNotNaN(self) -> "Column":  return self._wrap(~self._expr.is_nan())

    # ------------------------------------------------------------------
    # String operations
    # ------------------------------------------------------------------

    def contains(self, other) -> "Column":
        if isinstance(other, Column):
            # Use DuckDB-style; fall back to literal
            other = other._expr
        return self._wrap(self._expr.str.contains(str(other) if not isinstance(other, pl.Expr) else other))

    def startswith(self, other) -> "Column":
        return self._wrap(self._expr.str.starts_with(_lit_str(other)))

    def endswith(self, other) -> "Column":
        return self._wrap(self._expr.str.ends_with(_lit_str(other)))

    def like(self, pattern: str) -> "Column":
        import re
        regex = "^" + re.escape(pattern).replace(r"\%", ".*").replace(r"\_", ".") + "$"
        return self._wrap(self._expr.str.contains(regex))

    def rlike(self, pattern: str) -> "Column":
        return self._wrap(self._expr.str.contains(pattern))

    def substr(self, startPos, length=None) -> "Column":
        if isinstance(startPos, Column):
            startPos = 0  # best effort
        if isinstance(length, Column):
            length = None
        return self._wrap(self._expr.str.slice(startPos, length))

    # ------------------------------------------------------------------
    # Casting
    # ------------------------------------------------------------------

    def cast(self, dataType) -> "Column":
        from singlespark.sql.types import DataType, spark_type_to_polars, _parse_type_string
        if isinstance(dataType, str):
            dataType = _parse_type_string(dataType)
        if isinstance(dataType, DataType):
            polars_type = spark_type_to_polars(dataType)
        else:
            polars_type = dataType  # already a Polars type
        return self._wrap(self._expr.cast(polars_type))

    def astype(self, dataType) -> "Column":
        return self.cast(dataType)

    # ------------------------------------------------------------------
    # Aliasing / naming
    # ------------------------------------------------------------------

    def alias(self, *names: str) -> "Column":
        return self._wrap(self._expr.alias(names[0]))

    def name(self, name: str) -> "Column":
        return self.alias(name)

    # ------------------------------------------------------------------
    # Sorting markers
    # ------------------------------------------------------------------

    def asc(self) -> "Column":
        return _SortedColumn(self._expr, descending=False)

    def asc_nulls_first(self) -> "Column":
        return _SortedColumn(self._expr, descending=False, nulls_last=False)

    def asc_nulls_last(self) -> "Column":
        return _SortedColumn(self._expr, descending=False, nulls_last=True)

    def desc(self) -> "Column":
        return _SortedColumn(self._expr, descending=True)

    def desc_nulls_first(self) -> "Column":
        return _SortedColumn(self._expr, descending=True, nulls_last=False)

    def desc_nulls_last(self) -> "Column":
        return _SortedColumn(self._expr, descending=True, nulls_last=True)

    # ------------------------------------------------------------------
    # Membership / range
    # ------------------------------------------------------------------

    def isin(self, *values) -> "Column":
        if len(values) == 1 and isinstance(values[0], (list, tuple)):
            values = values[0]
        return self._wrap(self._expr.is_in(list(values)))

    def between(self, lowerBound, upperBound) -> "Column":
        return self._wrap(self._expr.is_between(_to_expr(lowerBound), _to_expr(upperBound)))

    # ------------------------------------------------------------------
    # when / otherwise (instance method chaining — rare but valid in PySpark)
    # ------------------------------------------------------------------

    def when(self, condition: "Column", value) -> "Column":
        raise NotImplementedError("Use F.when(condition, value) instead of col.when()")

    def otherwise(self, value) -> "Column":
        raise NotImplementedError("Use F.when(...).otherwise(value)")

    # ------------------------------------------------------------------
    # Item / field access
    # ------------------------------------------------------------------

    def __getitem__(self, key) -> "Column":
        if isinstance(key, (int,)):
            return self._wrap(self._expr.list.get(key))
        return self._wrap(self._expr.struct.field(str(key)))

    def getItem(self, key) -> "Column":
        return self.__getitem__(key)

    def getField(self, name: str) -> "Column":
        return self._wrap(self._expr.struct.field(name))

    # ------------------------------------------------------------------
    # Aggregation shortcuts (for use in agg())
    # ------------------------------------------------------------------

    def count(self) -> "Column":  return self._wrap(self._expr.count())
    def sum(self) -> "Column":    return self._wrap(self._expr.sum())
    def avg(self) -> "Column":    return self._wrap(self._expr.mean())
    def mean(self) -> "Column":   return self._wrap(self._expr.mean())
    def min(self) -> "Column":    return self._wrap(self._expr.min())
    def max(self) -> "Column":    return self._wrap(self._expr.max())
    def first(self) -> "Column":  return self._wrap(self._expr.first())
    def last(self) -> "Column":   return self._wrap(self._expr.last())

    def collect_list(self) -> "Column":  return self._wrap(self._expr.implode())
    def collect_set(self) -> "Column":   return self._wrap(self._expr.unique().implode())

    # ------------------------------------------------------------------
    # Misc
    # ------------------------------------------------------------------

    def bitwiseAND(self, other) -> "Column": return self._wrap(self._expr & _to_expr(other))
    def bitwiseOR(self, other) -> "Column":  return self._wrap(self._expr | _to_expr(other))
    def bitwiseXOR(self, other) -> "Column": return self._wrap(self._expr ^ _to_expr(other))

    def __repr__(self):
        return f"Column({self._expr})"


class _SortedColumn(Column):
    """Column with sort direction metadata (used by DataFrame.sort/orderBy)."""
    def __init__(self, expr: pl.Expr, descending: bool = False, nulls_last: bool | None = None):
        super().__init__(expr)
        self._descending = descending
        self._nulls_last = nulls_last


def _lit_str(val) -> str:
    if isinstance(val, Column):
        # best effort: extract literal string if possible
        return str(val._expr)
    return str(val)
