"""
singlespark.sql.functions — mirrors pyspark.sql.functions.

Import as:
    from singlespark.sql import functions as F
    F.col("name"), F.lit(42), F.when(F.col("age") > 18, "adult").otherwise("minor")
"""
from __future__ import annotations

from typing import Any

import polars as pl

from singlespark.column import Column, _to_expr

# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------

def _col(c) -> pl.Expr:
    """Accept Column, str, or pl.Expr, return pl.Expr."""
    if isinstance(c, Column):
        return c._expr
    if isinstance(c, str):
        return pl.col(c)
    if isinstance(c, pl.Expr):
        return c
    return pl.lit(c)


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


# ---------------------------------------------------------------------------
# Column reference + literal
# ---------------------------------------------------------------------------

def col(name: str) -> Column:
    """Reference a DataFrame column by name."""
    if name == "*":
        return _wrap(pl.all())
    return _wrap(pl.col(name))


def column(name: str) -> Column:
    return col(name)


def lit(value: Any) -> Column:
    """Create a Column of a literal value."""
    return _wrap(pl.lit(value))


# ---------------------------------------------------------------------------
# Aggregate functions
# ---------------------------------------------------------------------------

def count(col_name=None) -> Column:
    if col_name is None or col_name == "*":
        return _wrap(pl.len())
    return _wrap(_col(col_name).count())


def countDistinct(col_name, *other_cols) -> Column:
    expr = _col(col_name)
    return _wrap(expr.n_unique())


def approx_count_distinct(col_name, rsd: float = 0.05) -> Column:
    return countDistinct(col_name)


def sum(col_name) -> Column:           # noqa: A001
    return _wrap(_col(col_name).sum())


def sumDistinct(col_name) -> Column:
    return _wrap(_col(col_name).drop_nulls().unique().sum())


def avg(col_name) -> Column:
    return _wrap(_col(col_name).mean())


mean = avg


def min(col_name) -> Column:           # noqa: A001
    return _wrap(_col(col_name).min())


def max(col_name) -> Column:           # noqa: A001
    return _wrap(_col(col_name).max())


def first(col_name, ignorenulls: bool = False) -> Column:
    expr = _col(col_name)
    if ignorenulls:
        expr = expr.drop_nulls()
    return _wrap(expr.first())


def last(col_name, ignorenulls: bool = False) -> Column:
    expr = _col(col_name)
    if ignorenulls:
        expr = expr.drop_nulls()
    return _wrap(expr.last())


def collect_list(col_name) -> Column:
    return _wrap(_col(col_name).implode())


def collect_set(col_name) -> Column:
    return _wrap(_col(col_name).unique().implode())


def stddev(col_name) -> Column:
    return _wrap(_col(col_name).std())


stddev_samp = stddev


def stddev_pop(col_name) -> Column:
    return _wrap(_col(col_name).std(ddof=0))


def variance(col_name) -> Column:
    return _wrap(_col(col_name).var())


var_samp = variance


def var_pop(col_name) -> Column:
    return _wrap(_col(col_name).var(ddof=0))


def skewness(col_name) -> Column:
    return _wrap(_col(col_name).skew())


def kurtosis(col_name) -> Column:
    return _wrap(_col(col_name).kurtosis())


def corr(col1, col2) -> Column:
    return _wrap(pl.corr(_col(col1), _col(col2)))


def covar_samp(col1, col2) -> Column:
    return _wrap(pl.cov(_col(col1), _col(col2)))


def covar_pop(col1, col2) -> Column:
    return _wrap(pl.cov(_col(col1), _col(col2), ddof=0))


# ---------------------------------------------------------------------------
# Math functions
# ---------------------------------------------------------------------------

def abs(col_name) -> Column:           # noqa: A001
    return _wrap(_col(col_name).abs())


def sqrt(col_name) -> Column:
    return _wrap(_col(col_name).sqrt())


def cbrt(col_name) -> Column:
    return _wrap(_col(col_name) ** (1.0 / 3.0))


def exp(col_name) -> Column:
    return _wrap(_col(col_name).exp())


def log(col_name, base: float | None = None) -> Column:
    expr = _col(col_name).log(base) if base is not None else _col(col_name).log(2.718281828)
    return _wrap(expr)


def log2(col_name) -> Column:
    return _wrap(_col(col_name).log(2.0))


def log10(col_name) -> Column:
    return _wrap(_col(col_name).log(10.0))


def log1p(col_name) -> Column:
    return _wrap((_col(col_name) + pl.lit(1)).log(2.718281828))


def pow(col1, col2) -> Column:         # noqa: A001
    return _wrap(_col(col1) ** _col(col2))


def round(col_name, scale: int = 0) -> Column:  # noqa: A001
    return _wrap(_col(col_name).round(scale))


def bround(col_name, scale: int = 0) -> Column:
    return round(col_name, scale)


def ceil(col_name) -> Column:
    return _wrap(_col(col_name).ceil())


def floor(col_name) -> Column:
    return _wrap(_col(col_name).floor())


def rint(col_name) -> Column:
    return round(col_name, 0)


def signum(col_name) -> Column:
    return _wrap(_col(col_name).sign())


def factorial(col_name) -> Column:
    # Polars doesn't have factorial; compute via DuckDB hint or raise
    raise NotImplementedError("factorial() not supported — use selectExpr('factorial(col)')")


def greatest(*cols) -> Column:
    exprs = [_col(c) for c in cols]
    return _wrap(pl.max_horizontal(*exprs))


def least(*cols) -> Column:
    exprs = [_col(c) for c in cols]
    return _wrap(pl.min_horizontal(*exprs))


# Trig
def sin(col_name) -> Column:   return _wrap(_col(col_name).sin())
def cos(col_name) -> Column:   return _wrap(_col(col_name).cos())
def tan(col_name) -> Column:   return _wrap(_col(col_name).tan())
def asin(col_name) -> Column:  return _wrap(_col(col_name).arcsin())
def acos(col_name) -> Column:  return _wrap(_col(col_name).arccos())
def atan(col_name) -> Column:  return _wrap(_col(col_name).arctan())
def atan2(y, x) -> Column:
    return _wrap(pl.arctan2(_col(y), _col(x)))
def degrees(col_name) -> Column: return _wrap(_col(col_name) * pl.lit(180.0 / 3.141592653589793))
def radians(col_name) -> Column: return _wrap(_col(col_name) * pl.lit(3.141592653589793 / 180.0))
def hypot(col1, col2) -> Column:
    return _wrap((_col(col1)**2 + _col(col2)**2).sqrt())


def pi() -> Column:
    return lit(3.141592653589793)


# ---------------------------------------------------------------------------
# String functions
# ---------------------------------------------------------------------------

def lower(col_name) -> Column:
    return _wrap(_col(col_name).str.to_lowercase())


def upper(col_name) -> Column:
    return _wrap(_col(col_name).str.to_uppercase())


def length(col_name) -> Column:
    return _wrap(_col(col_name).str.len_chars())


def char_length(col_name) -> Column:
    return length(col_name)


def bit_length(col_name) -> Column:
    return _wrap(_col(col_name).str.len_bytes() * pl.lit(8))


def octet_length(col_name) -> Column:
    return _wrap(_col(col_name).str.len_bytes())


def trim(col_name) -> Column:
    return _wrap(_col(col_name).str.strip_chars())


def ltrim(col_name) -> Column:
    return _wrap(_col(col_name).str.strip_chars_start())


def rtrim(col_name) -> Column:
    return _wrap(_col(col_name).str.strip_chars_end())


def lpad(col_name, len: int, pad: str = " ") -> Column:
    return _wrap(_col(col_name).str.pad_start(len, pad[0] if pad else " "))


def rpad(col_name, len: int, pad: str = " ") -> Column:
    return _wrap(_col(col_name).str.pad_end(len, pad[0] if pad else " "))


def repeat(col_name, n) -> Column:
    return _wrap(_col(col_name).str.repeat(_col(n) if isinstance(n, Column) else n))


def reverse(col_name) -> Column:
    return _wrap(_col(col_name).str.reverse())


def substring(col_name, pos: int, len: int) -> Column:
    return _wrap(_col(col_name).str.slice(pos, len))


def substring_index(col_name, delim: str, count: int) -> Column:
    # Split and rejoin first/last count parts
    if count > 0:
        return _wrap(_col(col_name).str.splitn(delim, count + 1).list.slice(0, count).list.join(delim))
    else:
        return _wrap(_col(col_name).str.split(delim).list.slice(count).list.join(delim))


def split(col_name, pattern: str, limit: int = -1) -> Column:
    import re as _re
    # Polars str.split() only supports literal separators.
    # Detect regex-only patterns (e.g. r"\s+") and fall back to str.extract_all negation.
    _is_regex = bool(_re.search(r"[\[\]\\^$.*+?{}()|]", pattern))
    if _is_regex:
        # Build token pattern: match non-separator runs
        try:
            negated = f"[^{pattern.lstrip('[^').rstrip(']+')}]+"
            return _wrap(_col(col_name).str.extract_all(negated))
        except Exception:
            # Last resort: split on single space
            return _wrap(_col(col_name).str.split(" "))
    if limit > 0:
        return _wrap(_col(col_name).str.splitn(pattern, limit))
    return _wrap(_col(col_name).str.split(pattern))


def regexp_replace(col_name, pattern: str, replacement: str) -> Column:
    return _wrap(_col(col_name).str.replace_all(pattern, replacement))


def regexp_extract(col_name, pattern: str, idx: int = 0) -> Column:
    return _wrap(_col(col_name).str.extract(pattern, idx))


def regexp_extract_all(col_name, pattern: str) -> Column:
    return _wrap(_col(col_name).str.extract_all(pattern))


def instr(col_name, substr: str) -> Column:
    return _wrap(_col(col_name).str.find(substr))


def locate(substr: str, col_name, pos: int = 0) -> Column:
    return _wrap(_col(col_name).str.find(substr))


def concat(*cols) -> Column:
    exprs = [_col(c) for c in cols]
    return _wrap(pl.concat_str(*exprs, ignore_nulls=True))


def concat_ws(sep: str, *cols) -> Column:
    exprs = [_col(c) for c in cols]
    return _wrap(pl.concat_str(*exprs, separator=sep, ignore_nulls=True))


def format_string(format: str, *cols) -> Column:
    # Best effort via DuckDB printf-style — raise for now
    raise NotImplementedError("format_string() — use selectExpr(\"printf(fmt, col)\") instead")


def initcap(col_name) -> Column:
    return _wrap(_col(col_name).str.to_titlecase())


def soundex(col_name) -> Column:
    raise NotImplementedError("soundex() not supported in singlespark")


def levenshtein(left, right) -> Column:
    raise NotImplementedError("levenshtein() not supported — use selectExpr('levenshtein(a,b)')")


def translate(col_name, matching: str, replace: str) -> Column:
    raise NotImplementedError("translate() not supported in singlespark")


def encode(col_name, charset: str) -> Column:
    return _wrap(_col(col_name).str.encode("utf8"))


def decode(col_name, charset: str) -> Column:
    return _wrap(_col(col_name).str.decode("utf8"))


def base64(col_name) -> Column:
    return _wrap(_col(col_name).str.encode("base64"))


def unbase64(col_name) -> Column:
    return _wrap(_col(col_name).str.decode("base64"))


def ascii(col_name) -> Column:
    return _wrap(_col(col_name).str.slice(0, 1).cast(pl.Binary).bin.decode("utf-8"))


def chr(col_name) -> Column:
    raise NotImplementedError("chr() not supported")


def unhex(col_name) -> Column:
    return _wrap(_col(col_name).str.decode("hex"))


def hex(col_name) -> Column:      # noqa: A001
    return _wrap(_col(col_name).cast(pl.String))


# ---------------------------------------------------------------------------
# Null / conditional functions
# ---------------------------------------------------------------------------

def isnull(col_name) -> Column:
    return _wrap(_col(col_name).is_null())


def isnan(col_name) -> Column:
    return _wrap(_col(col_name).is_nan())


def coalesce(*cols) -> Column:
    exprs = [_col(c) for c in cols]
    return _wrap(pl.coalesce(exprs))


def nullif(col1, col2) -> Column:
    return _wrap(
        pl.when(_col(col1) == _col(col2)).then(pl.lit(None)).otherwise(_col(col1))
    )


def nvl(col1, col2) -> Column:
    return coalesce(col1, col2)


def nvl2(col1, col2, col3) -> Column:
    return _wrap(pl.when(_col(col1).is_not_null()).then(_col(col2)).otherwise(_col(col3)))


def ifnull(col1, col2) -> Column:
    return coalesce(col1, col2)


def nanvl(col1, col2) -> Column:
    return _wrap(pl.when(_col(col1).is_nan()).then(_col(col2)).otherwise(_col(col1)))


# ---------------------------------------------------------------------------
# when / otherwise
# ---------------------------------------------------------------------------

class _WhenBuilder:
    """Fluent builder for CASE-WHEN expressions."""

    def __init__(self, expr: pl.Expr):
        # expr is the partial pl.when(...).then(...) chain
        self._expr = expr

    def when(self, condition, value) -> "_WhenBuilder":
        cond_expr = _col(condition)
        val_expr = _to_expr(value)
        return _WhenBuilder(self._expr.when(cond_expr).then(val_expr))

    def otherwise(self, value) -> Column:
        return _wrap(self._expr.otherwise(_to_expr(value)))

    # Make _WhenBuilder usable as a Column in some contexts
    def alias(self, name: str) -> "_WhenBuilder":
        self._expr = self._expr.alias(name)
        return self

    def _as_column(self) -> Column:
        return _wrap(self._expr.otherwise(pl.lit(None)))

    # Allow wrapping in Column directly
    @property
    def _expr_col(self) -> Column:
        return _wrap(self._expr.otherwise(pl.lit(None)))


def when(condition, value) -> _WhenBuilder:
    """
    Evaluates a list of conditions and returns one of multiple possible result expressions.

    Example
    -------
        F.when(F.col("age") >= 18, "adult").when(F.col("age") >= 13, "teen").otherwise("child")
    """
    cond_expr = _col(condition)
    val_expr = _to_expr(value)
    return _WhenBuilder(pl.when(cond_expr).then(val_expr))


# Make _WhenBuilder behave like Column in DataFrame operations
_WhenBuilder.__class_getitem__ = lambda cls, _: cls  # type: ignore


# ---------------------------------------------------------------------------
# Date / time functions
# ---------------------------------------------------------------------------

def current_date() -> Column:
    return _wrap(pl.lit(pl.Series([pl.date(2000, 1, 1)])).first())  # placeholder


def current_timestamp() -> Column:
    import datetime
    return _wrap(pl.lit(datetime.datetime.now()))


def now() -> Column:
    return current_timestamp()


def year(col_name) -> Column:
    return _wrap(_col(col_name).dt.year())


def quarter(col_name) -> Column:
    return _wrap(_col(col_name).dt.quarter())


def month(col_name) -> Column:
    return _wrap(_col(col_name).dt.month())


def dayofmonth(col_name) -> Column:
    return _wrap(_col(col_name).dt.day())


def day(col_name) -> Column:
    return dayofmonth(col_name)


def dayofweek(col_name) -> Column:
    return _wrap(_col(col_name).dt.weekday() + pl.lit(1))


def dayofyear(col_name) -> Column:
    return _wrap(_col(col_name).dt.ordinal_day())


def weekofyear(col_name) -> Column:
    return _wrap(_col(col_name).dt.week())


def hour(col_name) -> Column:
    return _wrap(_col(col_name).dt.hour())


def minute(col_name) -> Column:
    return _wrap(_col(col_name).dt.minute())


def second(col_name) -> Column:
    return _wrap(_col(col_name).dt.second())


def unix_timestamp(col_name=None, format: str | None = None) -> Column:
    if col_name is None:
        import time
        return lit(int(time.time()))
    return _wrap(_col(col_name).dt.epoch(time_unit="s"))


def from_unixtime(col_name, format: str = "yyyy-MM-dd HH:mm:ss") -> Column:
    return _wrap(
        pl.from_epoch(_col(col_name), time_unit="s").dt.strftime(_java_fmt_to_python(format))
    )


def to_date(col_name, format: str | None = None) -> Column:
    if format:
        return _wrap(_col(col_name).str.strptime(pl.Date, _java_fmt_to_python(format), strict=False))
    return _wrap(_col(col_name).cast(pl.Date))


def to_timestamp(col_name, format: str | None = None) -> Column:
    if format:
        return _wrap(_col(col_name).str.strptime(pl.Datetime, _java_fmt_to_python(format), strict=False))
    return _wrap(_col(col_name).cast(pl.Datetime))


def date_format(col_name, format: str) -> Column:
    return _wrap(_col(col_name).dt.strftime(_java_fmt_to_python(format)))


def date_add(col_name, days) -> Column:
    return _wrap(_col(col_name) + pl.duration(days=_col(days) if isinstance(days, Column) else days))


def date_sub(col_name, days) -> Column:
    return _wrap(_col(col_name) - pl.duration(days=_col(days) if isinstance(days, Column) else days))


def datediff(end, start) -> Column:
    return _wrap((_col(end) - _col(start)).dt.total_days())


def months_between(date1, date2, roundOff: bool = True) -> Column:
    diff = (_col(date1) - _col(date2)).dt.total_days()
    result = diff / pl.lit(30.436875)
    if roundOff:
        result = result.round(8)
    return _wrap(result)


def add_months(col_name, months) -> Column:
    return _wrap(_col(col_name).dt.offset_by(f"{months}mo"))


def last_day(col_name) -> Column:
    return _wrap(_col(col_name).dt.month_end())


def next_day(col_name, dayOfWeek: str) -> Column:
    raise NotImplementedError("next_day() not yet supported")


def trunc(col_name, format: str) -> Column:
    fmt_map = {"year": "1y", "yyyy": "1y", "yy": "1y",
               "month": "1mo", "mon": "1mo", "mm": "1mo",
               "week": "1w", "day": "1d"}
    pl_fmt = fmt_map.get(format.lower(), "1d")
    return _wrap(_col(col_name).dt.truncate(pl_fmt))


def date_trunc(format: str, col_name) -> Column:
    return trunc(col_name, format)


# ---------------------------------------------------------------------------
# Array functions
# ---------------------------------------------------------------------------

def array(*cols) -> Column:
    exprs = [_col(c) for c in cols]
    return _wrap(pl.concat_list(*exprs))


def array_contains(col_name, value) -> Column:
    return _wrap(_col(col_name).list.contains(_to_expr(value)))


def array_distinct(col_name) -> Column:
    return _wrap(_col(col_name).list.unique())


def array_except(col1, col2) -> Column:
    raise NotImplementedError("array_except() not yet supported")


def array_intersect(col1, col2) -> Column:
    raise NotImplementedError("array_intersect() not yet supported")


def array_union(col1, col2) -> Column:
    return _wrap(pl.concat_list(_col(col1), _col(col2)).list.unique())


def array_join(col_name, delimiter: str, null_replacement: str | None = None) -> Column:
    if null_replacement:
        return _wrap(_col(col_name).list.join(delimiter))
    return _wrap(_col(col_name).list.join(delimiter))


def array_max(col_name) -> Column:
    return _wrap(_col(col_name).list.max())


def array_min(col_name) -> Column:
    return _wrap(_col(col_name).list.min())


def array_position(col_name, value) -> Column:
    raise NotImplementedError("array_position() not yet supported")


def array_remove(col_name, element) -> Column:
    raise NotImplementedError("array_remove() not yet supported")


def array_repeat(col_name, count) -> Column:
    raise NotImplementedError("array_repeat() not yet supported")


def array_sort(col_name) -> Column:
    return _wrap(_col(col_name).list.sort())


def array_zip(*cols) -> Column:
    exprs = [_col(c) for c in cols]
    return _wrap(pl.concat_list(*exprs))


def flatten(col_name) -> Column:
    return _wrap(_col(col_name).list.explode())


def size(col_name) -> Column:
    return _wrap(_col(col_name).list.len())


def explode(col_name) -> Column:
    return _wrap(_col(col_name).explode())


def explode_outer(col_name) -> Column:
    return explode(col_name)


def posexplode(col_name) -> Column:
    raise NotImplementedError("posexplode() not yet supported")


def posexplode_outer(col_name) -> Column:
    raise NotImplementedError("posexplode_outer() not yet supported")


def slice(col_name, start: int, length: int) -> Column:  # noqa: A001
    return _wrap(_col(col_name).list.slice(start, length))


def sort_array(col_name, asc: bool = True) -> Column:
    return _wrap(_col(col_name).list.sort(descending=not asc))


def reverse_array(col_name) -> Column:
    return _wrap(_col(col_name).list.reverse())


# ---------------------------------------------------------------------------
# Struct functions
# ---------------------------------------------------------------------------

def struct(*cols) -> Column:
    exprs = [_col(c) for c in cols]
    return _wrap(pl.struct(*exprs))


def create_map(*cols) -> Column:
    raise NotImplementedError("create_map() not yet supported")


def map_from_arrays(keys, values) -> Column:
    raise NotImplementedError("map_from_arrays() not yet supported")


def map_keys(col_name) -> Column:
    raise NotImplementedError("map_keys() not yet supported")


def map_values(col_name) -> Column:
    raise NotImplementedError("map_values() not yet supported")


# ---------------------------------------------------------------------------
# Window functions
# ---------------------------------------------------------------------------

def row_number() -> Column:
    return _wrap(pl.int_range(1, pl.len() + 1))


def rank() -> Column:
    return _wrap(pl.col("__rank__"))  # placeholder — requires Window spec


def dense_rank() -> Column:
    return _wrap(pl.col("__dense_rank__"))  # placeholder


def percent_rank() -> Column:
    raise NotImplementedError("percent_rank() requires Window specification")


def ntile(n: int) -> Column:
    raise NotImplementedError("ntile() requires Window specification")


def lag(col_name, offset: int = 1, default=None) -> Column:
    return _wrap(_col(col_name).shift(offset, fill_value=default))


def lead(col_name, offset: int = 1, default=None) -> Column:
    return _wrap(_col(col_name).shift(-offset, fill_value=default))


def cume_dist() -> Column:
    raise NotImplementedError("cume_dist() requires Window specification")


# ---------------------------------------------------------------------------
# Hash / ID functions
# ---------------------------------------------------------------------------

def hash(*cols) -> Column:      # noqa: A001
    return _wrap(pl.concat_str([_col(c).cast(pl.String) for c in cols]).hash())


def xxhash64(*cols) -> Column:
    return hash(*cols)


def md5(col_name) -> Column:
    raise NotImplementedError("md5() not supported in singlespark")


def sha1(col_name) -> Column:
    raise NotImplementedError("sha1() not supported in singlespark")


def sha2(col_name, numBits: int) -> Column:
    raise NotImplementedError("sha2() not supported in singlespark")


def crc32(col_name) -> Column:
    raise NotImplementedError("crc32() not supported in singlespark")


def monotonically_increasing_id() -> Column:
    return _wrap(pl.int_range(0, pl.len()))


def rand(seed: int | None = None) -> Column:
    if seed is not None:
        import numpy as np  # optional dep
        # Can't know N at expression build time; return expression
    return _wrap(pl.lit(0.5))  # best-effort placeholder


def randn(seed: int | None = None) -> Column:
    return _wrap(pl.lit(0.0))


def spark_partition_id() -> Column:
    return lit(0)


def input_file_name() -> Column:
    return lit("")


# ---------------------------------------------------------------------------
# Type casting
# ---------------------------------------------------------------------------

def cast(col_name, 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
    return _wrap(_col(col_name).cast(polars_type))


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

def broadcast(df):
    """No-op hint — returns df unchanged."""
    return df


def udf(f=None, returnType=None):
    """
    Create a user-defined function.

    Usage:
        double = udf(lambda x: x * 2, IntegerType())
        df.withColumn("doubled", double(col("val")))
    """
    from singlespark.sql.types import DataType

    def _make_udf(func):
        def _udf_col(*args):
            # Apply Python function to each row via Polars map
            arg_exprs = [_col(a) for a in args]

            def _apply(*row_vals):
                return func(*row_vals)

            # Use polars map_elements on the first column; multi-arg via struct
            if len(arg_exprs) == 1:
                result = arg_exprs[0].map_elements(_apply, return_dtype=pl.Utf8)
            else:
                combined = pl.struct(*arg_exprs)
                result = combined.map_elements(
                    lambda row: func(*row.values()), return_dtype=pl.Utf8
                )
            return _wrap(result)

        return _udf_col

    if f is None:
        return _make_udf
    return _make_udf(f)


def pandas_udf(f=None, returnType=None, functionType=None):
    """Simplified pandas_udf — falls back to plain udf behaviour."""
    return udf(f, returnType)


def expr(sql_expr: str) -> Column:
    """
    Parse a SQL expression string.
    The actual evaluation happens via DuckDB when the column is used in a DataFrame op.
    We store the raw SQL string in a special marker for DuckDB fallback.
    """
    return _wrap(pl.Expr.deserialize(sql_expr.encode(), format="json")
                 if False else pl.lit(sql_expr))  # deferred — use selectExpr instead


def element_at(col_name, extraction) -> Column:
    if isinstance(extraction, int):
        return _wrap(_col(col_name).list.get(extraction - 1))  # 1-indexed in Spark
    return _wrap(_col(col_name).struct.field(str(extraction)))


def get_json_object(col_name, path: str) -> Column:
    # path like "$.key.subkey"
    key = path.lstrip("$.")
    return _wrap(_col(col_name).str.json_path_match(f"$.{key}"))


def json_tuple(col_name, *fields: str) -> Column:
    # Returns multiple columns — not directly expressible as a single Column
    raise NotImplementedError("json_tuple() returns multiple columns; use selectExpr() instead")


def from_json(col_name, schema, options=None) -> Column:
    raise NotImplementedError("from_json() not yet supported — use selectExpr()")


def to_json(col_name, options=None) -> Column:
    return _wrap(_col(col_name).cast(pl.String))


def schema_of_json(col_name) -> Column:
    raise NotImplementedError("schema_of_json() not yet supported")


# ---------------------------------------------------------------------------
# Sorting helpers (standalone, for use in orderBy)
# ---------------------------------------------------------------------------

def asc(col_name) -> Column:
    from singlespark.column import _SortedColumn
    return _SortedColumn(pl.col(col_name) if isinstance(col_name, str) else _col(col_name),
                         descending=False)


def desc(col_name) -> Column:
    from singlespark.column import _SortedColumn
    return _SortedColumn(pl.col(col_name) if isinstance(col_name, str) else _col(col_name),
                         descending=True)


def asc_nulls_first(col_name) -> Column:
    from singlespark.column import _SortedColumn
    return _SortedColumn(pl.col(col_name) if isinstance(col_name, str) else _col(col_name),
                         descending=False, nulls_last=False)


def asc_nulls_last(col_name) -> Column:
    from singlespark.column import _SortedColumn
    return _SortedColumn(pl.col(col_name) if isinstance(col_name, str) else _col(col_name),
                         descending=False, nulls_last=True)


def desc_nulls_first(col_name) -> Column:
    from singlespark.column import _SortedColumn
    return _SortedColumn(pl.col(col_name) if isinstance(col_name, str) else _col(col_name),
                         descending=True, nulls_last=False)


def desc_nulls_last(col_name) -> Column:
    from singlespark.column import _SortedColumn
    return _SortedColumn(pl.col(col_name) if isinstance(col_name, str) else _col(col_name),
                         descending=True, nulls_last=True)


# ---------------------------------------------------------------------------
# Internal helper: Java datetime format → Python strftime
# ---------------------------------------------------------------------------

def _java_fmt_to_python(fmt: str) -> str:
    """Convert a Java SimpleDateFormat pattern to Python strftime."""
    replacements = [
        ("yyyy", "%Y"),
        ("yy", "%y"),
        ("MMMM", "%B"),
        ("MMM", "%b"),
        ("MM", "%m"),
        ("M", "%-m"),
        ("dd", "%d"),
        ("d", "%-d"),
        ("HH", "%H"),
        ("H", "%-H"),
        ("hh", "%I"),
        ("h", "%-I"),
        ("mm", "%M"),
        ("ss", "%S"),
        ("SSS", "%f"),
        ("a", "%p"),
        ("EEE", "%a"),
        ("EEEE", "%A"),
    ]
    result = fmt
    for java_pat, py_pat in replacements:
        result = result.replace(java_pat, py_pat)
    return result
