# Contains code from https://github.com/ibis-project/ibis/blob/main/ibis/backends/sql/compilers/bigquery/__init__.py
"""Module to convert from Ibis expression to SQL string."""

from __future__ import annotations

import datetime
import decimal
import math
import re
from typing import Any, TYPE_CHECKING

from bigframes_vendored.ibis import util
import bigframes_vendored.ibis.backends.bigquery.datatypes as bq_datatypes
from bigframes_vendored.ibis.backends.sql.compilers.base import (
    AggGen,
    NULL,
    SQLGlotCompiler,
    STAR,
)
from bigframes_vendored.ibis.backends.sql.datatypes import BigQueryType, BigQueryUDFType
from bigframes_vendored.ibis.backends.sql.rewrites import (
    exclude_unsupported_window_frame_from_ops,
    exclude_unsupported_window_frame_from_rank,
    exclude_unsupported_window_frame_from_row_number,
)
import bigframes_vendored.ibis.common.exceptions as ibis_exceptions
from bigframes_vendored.ibis.common.temporal import (
    DateUnit,
    IntervalUnit,
    TimestampUnit,
    TimeUnit,
)
import bigframes_vendored.ibis.expr.datatypes as dt
import bigframes_vendored.ibis.expr.operations as ops
import numpy as np
import sqlglot as sg
from sqlglot.dialects import BigQuery
import sqlglot.expressions as sge

if TYPE_CHECKING:
    from collections.abc import Mapping

    import bigframes_vendored.ibis.expr.types as ir


_NAME_REGEX = re.compile(r'[^!"$()*,./;?@[\\\]^`{}~\n]+')


_MEMTABLE_PATTERN = re.compile(
    r"^_?ibis_(?:[A-Za-z_][A-Za-z_0-9]*)_memtable_[a-z0-9]{26}$"
)


def _qualify_memtable(
    node: sge.Expression, *, dataset: str | None, project: str | None
) -> sge.Expression:
    """Add a BigQuery dataset and project to memtable references."""
    if isinstance(node, sge.Table) and _MEMTABLE_PATTERN.match(node.name) is not None:
        node.args["db"] = dataset
        node.args["catalog"] = project
    return node


def _remove_null_ordering_from_unsupported_window(
    node: sge.Expression,
) -> sge.Expression:
    """Remove null ordering in window frame clauses not supported by BigQuery.
    BigQuery has only partial support for NULL FIRST/LAST in RANGE windows so
    we remove it from any window frame clause that doesn't support it.
    Here's the support matrix:
    ✅ sum(x) over (order by y desc nulls last)
    🚫 sum(x) over (order by y asc nulls last)
    ✅ sum(x) over (order by y asc nulls first)
    🚫 sum(x) over (order by y desc nulls first)
    """
    if isinstance(node, sge.Window):
        order = node.args.get("order")
        if order is not None:
            for key in order.args["expressions"]:
                kargs = key.args
                if kargs.get("desc") is True and kargs.get("nulls_first", False):
                    kargs["nulls_first"] = False
                elif kargs.get("desc") is False and not kargs.setdefault(
                    "nulls_first", True
                ):
                    kargs["nulls_first"] = True
    return node


class BigQueryCompiler(SQLGlotCompiler):
    dialect = BigQuery
    type_mapper = BigQueryType
    udf_type_mapper = BigQueryUDFType

    agg = AggGen(supports_order_by=True)

    rewrites = (
        exclude_unsupported_window_frame_from_ops,
        exclude_unsupported_window_frame_from_row_number,
        exclude_unsupported_window_frame_from_rank,
        *SQLGlotCompiler.rewrites,
    )

    supports_qualify = True

    UNSUPPORTED_OPS = (
        ops.DateDiff,
        ops.ExtractAuthority,
        ops.ExtractUserInfo,
        ops.FindInSet,
        ops.Median,
        ops.RegexSplit,
        ops.RowID,
        ops.TimestampDiff,
    )

    NAN = sge.Cast(
        this=sge.convert("NaN"), to=sge.DataType(this=sge.DataType.Type.DOUBLE)
    )
    POS_INF = sge.Cast(
        this=sge.convert("Infinity"), to=sge.DataType(this=sge.DataType.Type.DOUBLE)
    )
    NEG_INF = sge.Cast(
        this=sge.convert("-Infinity"), to=sge.DataType(this=sge.DataType.Type.DOUBLE)
    )

    SIMPLE_OPS = {
        ops.Arbitrary: "any_value",
        ops.StringAscii: "ascii",
        ops.BitAnd: "bit_and",
        ops.BitOr: "bit_or",
        ops.BitXor: "bit_xor",
        ops.DateFromYMD: "date",
        ops.Divide: "ieee_divide",
        ops.EndsWith: "ends_with",
        ops.GeoArea: "st_area",
        ops.GeoAsBinary: "st_asbinary",
        ops.GeoAsText: "st_astext",
        ops.GeoAzimuth: "st_azimuth",
        ops.GeoBuffer: "st_buffer",
        ops.GeoCentroid: "st_centroid",
        ops.GeoContains: "st_contains",
        ops.GeoCoveredBy: "st_coveredby",
        ops.GeoCovers: "st_covers",
        ops.GeoDWithin: "st_dwithin",
        ops.GeoDifference: "st_difference",
        ops.GeoDisjoint: "st_disjoint",
        ops.GeoDistance: "st_distance",
        ops.GeoEndPoint: "st_endpoint",
        ops.GeoEquals: "st_equals",
        ops.GeoGeometryType: "st_geometrytype",
        ops.GeoIntersection: "st_intersection",
        ops.GeoIntersects: "st_intersects",
        ops.GeoLength: "st_length",
        ops.GeoMaxDistance: "st_maxdistance",
        ops.GeoNPoints: "st_numpoints",
        ops.GeoPerimeter: "st_perimeter",
        ops.GeoPoint: "st_geogpoint",
        ops.GeoPointN: "st_pointn",
        ops.GeoStartPoint: "st_startpoint",
        ops.GeoTouches: "st_touches",
        ops.GeoUnaryUnion: "st_union_agg",
        ops.GeoUnion: "st_union",
        ops.GeoWithin: "st_within",
        ops.GeoX: "st_x",
        ops.GeoY: "st_y",
        ops.Hash: "farm_fingerprint",
        ops.IsInf: "is_inf",
        ops.IsNan: "is_nan",
        ops.Log10: "log10",
        ops.LPad: "lpad",
        ops.RPad: "rpad",
        ops.Levenshtein: "edit_distance",
        ops.Modulus: "mod",
        ops.RegexReplace: "regexp_replace",
        ops.RegexSearch: "regexp_contains",
        ops.Time: "time",
        ops.TimeFromHMS: "time_from_parts",
        ops.TimestampNow: "current_timestamp",
        ops.ExtractHost: "net.host",
    }

    def to_sqlglot(
        self,
        expr: ir.Expr,
        *,
        limit: str | None = None,
        params: Mapping[ir.Expr, Any] | None = None,
        session_dataset_id: str | None = None,
        session_project: str | None = None,
    ) -> Any:
        """Compile an Ibis expression.
        Parameters
        ----------
        expr
            Ibis expression
        limit
            For expressions yielding result sets; retrieve at most this number
            of values/rows. Overrides any limit already set on the expression.
        params
            Named unbound parameters
        session_dataset_id
            Optional dataset ID to qualify memtable references.
        session_project
            Optional project ID to qualify memtable references.
        Returns
        -------
        Any
            The output of compilation. The type of this value depends on the
            backend.
        """
        sql = super().to_sqlglot(expr, limit=limit, params=params)

        table_expr = expr.as_table()
        geocols = getattr(table_expr.schema(), "geospatial", None)

        result = sql.transform(
            _qualify_memtable,
            dataset=session_dataset_id,
            project=session_project,
        ).transform(_remove_null_ordering_from_unsupported_window)

        if geocols:
            # if there are any geospatial columns, we have to convert them to WKB,
            # so interactive mode knows how to display them
            #
            # by default bigquery returns data to python as WKT, and there's really
            # no point in supporting both if we don't need to.
            quoted = self.quoted
            result = sg.select(
                sge.Star(
                    replace=[
                        self.f.st_asbinary(sg.column(col, quoted=quoted)).as_(
                            col, quoted=quoted
                        )
                        for col in geocols
                    ]
                )
            ).from_(result.subquery())

        sources = []

        for udf_node in table_expr.op().find(ops.ScalarUDF):
            compile_func = getattr(
                self, f"_compile_{udf_node.__input_type__.name.lower()}_udf"
            )
            if sql := compile_func(udf_node):
                sources.append(sql)

        if not sources:
            return result

        sources.append(result)
        return sources

    def visit_BoundingBox(self, op, *, arg):
        name = type(op).__name__[len("Geo") :].lower()
        return sge.Dot(
            this=self.f.st_boundingbox(arg), expression=sg.to_identifier(name)
        )

    visit_GeoXMax = visit_GeoXMin = visit_GeoYMax = visit_GeoYMin = visit_BoundingBox

    def visit_GeoRegionStats(self, op, *, arg, raster_id, band, include, options):
        args = [arg, raster_id]
        if op.band:
            args.append(sge.Kwarg(this="band", expression=band))
        if op.include:
            args.append(sge.Kwarg(this="include", expression=include))
        if op.options:
            args.append(sge.Kwarg(this="options", expression=options))
        return sge.func("ST_REGIONSTATS", *args)

    def visit_GeoSimplify(self, op, *, arg, tolerance, preserve_collapsed):
        if (
            not isinstance(op.preserve_collapsed, ops.Literal)
            or op.preserve_collapsed.value
        ):
            raise ibis_exceptions.UnsupportedOperationError(
                "BigQuery simplify does not support preserving collapsed geometries, "
                "pass preserve_collapsed=False"
            )
        return self.f.st_simplify(arg, tolerance)

    def visit_ApproxMedian(self, op, *, arg, where):
        return self.agg.approx_quantiles(arg, 2, where=where)[self.f.offset(1)]

    def visit_Pi(self, op):
        return self.f.acos(-1)

    def visit_E(self, op):
        return self.f.exp(1)

    def visit_TimeDelta(self, op, *, left, right, part):
        return self.f.time_diff(left, right, part, dialect=self.dialect)

    def visit_DateDelta(self, op, *, left, right, part):
        return self.f.date_diff(left, right, part, dialect=self.dialect)

    def visit_TimestampDelta(self, op, *, left, right, part):
        left_tz = op.left.dtype.timezone
        right_tz = op.right.dtype.timezone

        if left_tz is None and right_tz is None:
            return self.f.datetime_diff(left, right, part)
        elif left_tz is not None and right_tz is not None:
            return self.f.timestamp_diff(left, right, part)

        raise ibis_exceptions.UnsupportedOperationError(
            "timestamp difference with mixed timezone/timezoneless values is not implemented"
        )

    def visit_GroupConcat(self, op, *, arg, sep, where, order_by):
        if where is not None:
            arg = self.if_(where, arg, NULL)

        if order_by:
            sep = sge.Order(this=sep, expressions=order_by)

        return sge.GroupConcat(this=arg, separator=sep)

    def visit_ApproxQuantile(self, op, *, arg, quantile, where):
        if not isinstance(op.quantile, ops.Literal):
            raise ibis_exceptions.UnsupportedOperationError(
                "quantile must be a literal in BigQuery"
            )

        # BigQuery syntax is `APPROX_QUANTILES(col, resolution)` to return
        # `resolution + 1` quantiles array. To handle this, we compute the
        # resolution ourselves then restructure the output array as needed.
        # To avoid excessive resolution we arbitrarily cap it at 100,000 -
        # since these are approximate quantiles anyway this seems fine.
        quantiles = util.promote_list(op.quantile.value)
        fracs = [decimal.Decimal(str(q)).as_integer_ratio() for q in quantiles]
        resolution = min(math.lcm(*(den for _, den in fracs)), 100_000)
        indices = [(num * resolution) // den for num, den in fracs]

        if where is not None:
            arg = self.if_(where, arg, NULL)

        if not op.arg.dtype.is_floating():
            arg = self.cast(arg, dt.float64)

        array = self.f.approx_quantiles(
            arg, sge.IgnoreNulls(this=sge.convert(resolution))
        )
        if isinstance(op, ops.ApproxQuantile):
            return array[indices[0]]

        if indices == list(range(resolution + 1)):
            return array
        else:
            return sge.Array(expressions=[array[i] for i in indices])

    visit_ApproxMultiQuantile = visit_ApproxQuantile

    def visit_FloorDivide(self, op, *, left, right):
        return self.cast(self.f.floor(self.f.ieee_divide(left, right)), op.dtype)

    def visit_Log2(self, op, *, arg):
        return self.f.log(arg, 2, dialect=self.dialect)

    def visit_Log(self, op, *, arg, base):
        if base is None:
            return self.f.ln(arg)
        return self.f.log(arg, base, dialect=self.dialect)

    def visit_ArrayRepeat(self, op, *, arg, times):
        start = step = 1
        array_length = self.f.array_length(arg)
        stop = self.f.greatest(times, 0) * array_length
        i = sg.to_identifier("i")
        idx = self.f.coalesce(
            self.f.nullif(self.f.mod(i, array_length), 0), array_length
        )
        series = self.f.generate_array(start, stop, step)
        return self.f.array(
            sg.select(arg[self.f.safe_ordinal(idx)]).from_(self._unnest(series, as_=i))
        )

    def visit_NthValue(self, op, *, arg, nth):
        if not isinstance(op.nth, ops.Literal):
            raise ibis_exceptions.UnsupportedOperationError(
                f"BigQuery `nth` must be a literal; got {type(op.nth)}"
            )
        return self.f.nth_value(arg, nth)

    def visit_StrRight(self, op, *, arg, nchars):
        return self.f.substr(arg, -self.f.least(self.f.length(arg), nchars))

    def visit_StringJoin(self, op, *, arg, sep):
        return self.f.array_to_string(self.f.array(*arg), sep)

    def visit_DayOfWeekIndex(self, op, *, arg):
        return self.f.mod(self.f.extract(self.v.dayofweek, arg) + 5, 7)

    def visit_DayOfWeekName(self, op, *, arg):
        return self.f.initcap(sge.Cast(this=arg, to="STRING FORMAT 'DAY'"))

    def visit_StringToTimestamp(self, op, *, arg, format_str):
        if (timezone := op.dtype.timezone) is not None:
            return self.f.parse_timestamp(format_str, arg, timezone)
        return self.f.parse_datetime(format_str, arg)

    def visit_ArrayCollect(self, op, *, arg, where, order_by, include_null):
        if where is not None and include_null:
            raise ibis_exceptions.UnsupportedOperationError(
                "Combining `include_null=True` and `where` is not supported "
                "by bigquery"
            )
        out = self.agg.array_agg(arg, where=where, order_by=order_by)
        if not include_null:
            out = sge.IgnoreNulls(this=out)
        return out

    def _neg_idx_to_pos(self, arg, idx):
        return self.if_(idx < 0, self.f.array_length(arg) + idx, idx)

    def visit_ArraySlice(self, op, *, arg, start, stop):
        index = sg.to_identifier("bq_arr_slice")
        cond = [index >= self._neg_idx_to_pos(arg, start)]

        if stop is not None:
            cond.append(index < self._neg_idx_to_pos(arg, stop))

        el = sg.to_identifier("el")
        return self.f.array(
            sg.select(el).from_(self._unnest(arg, as_=el, offset=index)).where(*cond)
        )

    def visit_ArrayIndex(self, op, *, arg, index):
        return arg[self.f.safe_offset(index)]

    def visit_ArrayContains(self, op, *, arg, other):
        name = sg.to_identifier(util.gen_name("bq_arr_contains"))
        return sge.Exists(
            this=sg.select(sge.convert(1))
            .from_(self._unnest(arg, as_=name))
            .where(name.eq(other))
        )

    def visit_StringContains(self, op, *, haystack, needle):
        return self.f.strpos(haystack, needle) > 0

    def visit_StringFind(self, op, *, arg, substr, start, end):
        if start is not None:
            raise NotImplementedError(
                "`start` not implemented for BigQuery string find"
            )
        if end is not None:
            raise NotImplementedError("`end` not implemented for BigQuery string find")
        return self.f.strpos(arg, substr)

    def visit_TimestampFromYMDHMS(
        self, op, *, year, month, day, hours, minutes, seconds
    ):
        return self.f.anon.DATETIME(year, month, day, hours, minutes, seconds)

    def visit_NonNullLiteral(self, op, *, value, dtype):
        if dtype.is_inet() or dtype.is_macaddr():
            return sge.convert(str(value))
        elif dtype.is_timestamp():
            funcname = "DATETIME" if dtype.timezone is None else "TIMESTAMP"
            return self.f.anon[funcname](value.isoformat())
        elif dtype.is_date():
            return self.f.date_from_parts(value.year, value.month, value.day)
        elif dtype.is_time():
            time = self.f.time_from_parts(value.hour, value.minute, value.second)
            if micros := value.microsecond:
                # bigquery doesn't support `time(12, 34, 56.789101)`, AKA a
                # float seconds specifier, so add any non-zero micros to the
                # time value
                return sge.TimeAdd(
                    this=time, expression=sge.convert(micros), unit=self.v.MICROSECOND
                )
            return time
        elif dtype.is_binary():
            return sge.Cast(
                this=sge.convert(value.hex()),
                to=sge.DataType(this=sge.DataType.Type.BINARY),
                format=sge.convert("HEX"),
            )
        elif dtype.is_interval():
            if dtype.unit == IntervalUnit.NANOSECOND:
                raise ibis_exceptions.UnsupportedOperationError(
                    "BigQuery does not support nanosecond intervals"
                )
        elif dtype.is_uuid():
            return sge.convert(str(value))

        elif dtype.is_int64():
            # allows directly using values out of a duration arrow array
            if isinstance(value, datetime.timedelta):
                value = (
                    (value.days * 3600 * 24) + value.seconds
                ) * 1_000_000 + value.microseconds
            return sge.convert(np.int64(value))
        return None

    def visit_IntervalFromInteger(self, op, *, arg, unit):
        if unit == IntervalUnit.NANOSECOND:
            raise ibis_exceptions.UnsupportedOperationError(
                "BigQuery does not support nanosecond intervals"
            )
        return sge.Interval(this=arg, unit=self.v[unit.singular])

    def visit_Strftime(self, op, *, arg, format_str):
        arg_dtype = op.arg.dtype
        if arg_dtype.is_timestamp():
            if (timezone := arg_dtype.timezone) is None:
                return self.f.format_datetime(format_str, arg)
            else:
                return self.f.format_timestamp(format_str, arg, timezone)
        elif arg_dtype.is_date():
            return self.f.format_date(format_str, arg)
        else:
            assert arg_dtype.is_time(), arg_dtype
            return self.f.format_time(format_str, arg)

    def visit_IntervalMultiply(self, op, *, left, right):
        unit = self.v[op.left.dtype.resolution.upper()]
        return sge.Interval(this=self.f.extract(unit, left) * right, unit=unit)

    def visit_TimestampFromUNIX(self, op, *, arg, unit):
        unit = op.unit
        if unit == TimestampUnit.SECOND:
            return self.f.timestamp_seconds(arg)
        elif unit == TimestampUnit.MILLISECOND:
            return self.f.timestamp_millis(arg)
        elif unit == TimestampUnit.MICROSECOND:
            return self.f.timestamp_micros(arg)
        elif unit == TimestampUnit.NANOSECOND:
            return self.f.timestamp_micros(
                self.cast(self.f.round(arg / 1_000), dt.int64)
            )
        else:
            raise ibis_exceptions.UnsupportedOperationError(
                f"Unit not supported: {unit}"
            )

    def visit_Cast(self, op, *, arg, to):
        from_ = op.arg.dtype
        if from_.is_timestamp() and to.is_integer():
            return self.f.unix_micros(arg)
        elif from_.is_integer() and to.is_timestamp():
            return self.f.timestamp_seconds(arg)
        elif from_.is_interval() and to.is_integer():
            if from_.unit in {
                IntervalUnit.WEEK,
                IntervalUnit.QUARTER,
                IntervalUnit.NANOSECOND,
            }:
                raise ibis_exceptions.UnsupportedOperationError(
                    f"BigQuery does not allow extracting date part `{from_.unit}` from intervals"
                )
            return self.f.extract(self.v[to.resolution.upper()], arg)
        elif (from_.is_floating() or from_.is_decimal()) and to.is_integer():
            return self.cast(self.f.trunc(arg), dt.int64)
        return super().visit_Cast(op, arg=arg, to=to)

    def visit_JSONGetItem(self, op, *, arg, index):
        return arg[index]

    def visit_UnwrapJSONString(self, op, *, arg):
        return self.f.anon["safe.string"](arg)

    def visit_UnwrapJSONInt64(self, op, *, arg):
        return self.f.anon["safe.int64"](arg)

    def visit_UnwrapJSONFloat64(self, op, *, arg):
        return self.f.anon["safe.float64"](arg)

    def visit_UnwrapJSONBoolean(self, op, *, arg):
        return self.f.anon["safe.bool"](arg)

    def visit_ExtractEpochSeconds(self, op, *, arg):
        return self.f.unix_seconds(arg)

    def visit_ExtractWeekOfYear(self, op, *, arg):
        return self.f.extract(self.v.isoweek, arg)

    def visit_ExtractIsoYear(self, op, *, arg):
        return self.f.extract(self.v.isoyear, arg)

    def visit_ExtractMillisecond(self, op, *, arg):
        return self.f.extract(self.v.millisecond, arg)

    def visit_ExtractMicrosecond(self, op, *, arg):
        return self.f.extract(self.v.microsecond, arg)

    def visit_TimestampTruncate(self, op, *, arg, unit):
        if unit == IntervalUnit.NANOSECOND:
            raise ibis_exceptions.UnsupportedOperationError(
                f"BigQuery does not support truncating {op.arg.dtype} values to unit {unit!r}"
            )
        elif unit == IntervalUnit.WEEK:
            unit = "WEEK(MONDAY)"
        else:
            unit = unit.name
        return self.f.timestamp_trunc(arg, self.v[unit], dialect=self.dialect)

    def visit_DateTruncate(self, op, *, arg, unit):
        if unit == DateUnit.WEEK:
            unit = "WEEK(MONDAY)"
        else:
            unit = unit.name
        return self.f.date_trunc(arg, self.v[unit], dialect=self.dialect)

    def visit_TimeTruncate(self, op, *, arg, unit):
        if unit == TimeUnit.NANOSECOND:
            raise ibis_exceptions.UnsupportedOperationError(
                f"BigQuery does not support truncating {op.arg.dtype} values to unit {unit!r}"
            )
        else:
            unit = unit.name
        return self.f.time_trunc(arg, self.v[unit], dialect=self.dialect)

    def _nullifzero(self, step, zero, step_dtype):
        if step_dtype.is_interval():
            return self.if_(step.eq(zero), NULL, step)
        return self.f.nullif(step, zero)

    def _zero(self, dtype):
        if dtype.is_interval():
            return self.f.make_interval()
        return sge.convert(0)

    def _sign(self, value, dtype):
        if dtype.is_interval():
            zero = self._zero(dtype)
            return sge.Case(
                ifs=[
                    self.if_(value < zero, -1),
                    self.if_(value.eq(zero), 0),
                    self.if_(value > zero, 1),
                ],
                default=NULL,
            )
        return self.f.sign(value)

    def _make_range(self, func, start, stop, step, step_dtype):
        step_sign = self._sign(step, step_dtype)
        delta_sign = self._sign(stop - start, step_dtype)
        zero = self._zero(step_dtype)
        nullifzero = self._nullifzero(step, zero, step_dtype)
        condition = sg.and_(sg.not_(nullifzero.is_(NULL)), step_sign.eq(delta_sign))
        gen_array = func(start, stop, step)
        name = sg.to_identifier(util.gen_name("bq_arr_range"))
        inner = (
            sg.select(name)
            .from_(self._unnest(gen_array, as_=name))
            .where(name.neq(stop))
        )
        return self.if_(condition, self.f.array(inner), self.f.array())

    def visit_IntegerRange(self, op, *, start, stop, step):
        return self._make_range(self.f.generate_array, start, stop, step, op.step.dtype)

    def visit_TimestampRange(self, op, *, start, stop, step):
        if op.start.dtype.timezone is None or op.stop.dtype.timezone is None:
            raise ibis_exceptions.IbisTypeError(
                "Timestamps without timezone values are not supported when generating timestamp ranges"
            )
        return self._make_range(
            self.f.generate_timestamp_array, start, stop, step, op.step.dtype
        )

    def visit_First(self, op, *, arg, where, order_by, include_null):
        if where is not None:
            arg = self.if_(where, arg, NULL)
            if include_null:
                raise ibis_exceptions.UnsupportedOperationError(
                    "Combining `include_null=True` and `where` is not supported "
                    "by bigquery"
                )

        if order_by:
            arg = sge.Order(this=arg, expressions=order_by)

        if not include_null:
            arg = sge.IgnoreNulls(this=arg)

        array = self.f.array_agg(sge.Limit(this=arg, expression=sge.convert(1)))
        return array[self.f.safe_offset(0)]

    def visit_Last(self, op, *, arg, where, order_by, include_null):
        if where is not None:
            arg = self.if_(where, arg, NULL)
            if include_null:
                raise ibis_exceptions.UnsupportedOperationError(
                    "Combining `include_null=True` and `where` is not supported "
                    "by bigquery"
                )

        if order_by:
            arg = sge.Order(this=arg, expressions=order_by)

        if not include_null:
            arg = sge.IgnoreNulls(this=arg)

        array = self.f.array_reverse(self.f.array_agg(arg))
        return array[self.f.safe_offset(0)]

    def visit_ArrayFilter(self, op, *, arg, body, param):
        return self.f.array(
            sg.select(param).from_(self._unnest(arg, as_=param)).where(body)
        )

    def visit_ArrayMap(self, op, *, arg, body, param):
        return self.f.array(sg.select(body).from_(self._unnest(arg, as_=param)))

    def visit_ArrayReduce(self, op, *, arg, body, param):
        return sg.select(body).from_(self._unnest(arg, as_=param)).subquery()

    def visit_ArrayZip(self, op, *, arg):
        lengths = [self.f.array_length(arr) - 1 for arr in arg]
        idx = sg.to_identifier(util.gen_name("bq_arr_idx"))
        indices = self._unnest(
            self.f.generate_array(0, self.f.greatest(*lengths)), as_=idx
        )
        struct_fields = [
            arr[self.f.safe_offset(idx)].as_(name)
            for name, arr in zip(op.dtype.value_type.names, arg)
        ]
        return self.f.array(
            sge.Select(kind="STRUCT", expressions=struct_fields).from_(indices)
        )

    def visit_ArrayPosition(self, op, *, arg, other):
        name = sg.to_identifier(util.gen_name("bq_arr"))
        idx = sg.to_identifier(util.gen_name("bq_arr_idx"))
        unnest = self._unnest(arg, as_=name, offset=idx)
        return self.f.coalesce(
            sg.select(idx + 1).from_(unnest).where(name.eq(other)).limit(1).subquery(),
            0,
        )

    def _unnest(self, expression, *, as_, offset=None):
        alias = sge.TableAlias(columns=[sg.to_identifier(as_)])
        return sge.Unnest(expressions=[expression], alias=alias, offset=offset)

    def visit_ArrayRemove(self, op, *, arg, other):
        name = sg.to_identifier(util.gen_name("bq_arr"))
        unnest = self._unnest(arg, as_=name)
        both_null = sg.and_(name.is_(NULL), other.is_(NULL))
        cond = sg.or_(name.neq(other), both_null)
        return self.f.array(sg.select(name).from_(unnest).where(cond))

    def visit_ArrayDistinct(self, op, *, arg):
        name = util.gen_name("bq_arr")
        return self.f.array(
            sg.select(name).distinct().from_(self._unnest(arg, as_=name))
        )

    def visit_ArraySort(self, op, *, arg):
        name = util.gen_name("bq_arr")
        return self.f.array(
            sg.select(name).from_(self._unnest(arg, as_=name)).order_by(name)
        )

    def visit_ArrayUnion(self, op, *, left, right):
        lname = util.gen_name("bq_arr_left")
        rname = util.gen_name("bq_arr_right")
        lhs = sg.select(lname).from_(self._unnest(left, as_=lname))
        rhs = sg.select(rname).from_(self._unnest(right, as_=rname))
        return self.f.array(sg.union(lhs, rhs, distinct=True))

    def visit_ArrayIntersect(self, op, *, left, right):
        lname = util.gen_name("bq_arr_left")
        rname = util.gen_name("bq_arr_right")
        lhs = sg.select(lname).from_(self._unnest(left, as_=lname))
        rhs = sg.select(rname).from_(self._unnest(right, as_=rname))
        return self.f.array(sg.intersect(lhs, rhs, distinct=True))

    def visit_RegexExtract(self, op, *, arg, pattern, index):
        matches = self.f.regexp_contains(arg, pattern)
        nonzero_index_replace = self.f.regexp_replace(
            arg,
            self.f.concat(".*?", pattern, ".*"),
            self.f.concat("\\", self.cast(index, dt.string)),
        )
        zero_index_replace = self.f.regexp_replace(
            arg, self.f.concat(".*?", self.f.concat("(", pattern, ")"), ".*"), "\\1"
        )
        extract = self.if_(index.eq(0), zero_index_replace, nonzero_index_replace)
        return self.if_(matches, extract, NULL)

    def visit_TimestampAddSub(self, op, *, left, right):
        if not isinstance(right, sge.Interval):
            raise ibis_exceptions.OperationNotDefinedError(
                "BigQuery does not support non-literals on the right side of timestamp add/subtract"
            )
        if (unit := op.right.dtype.unit) == IntervalUnit.NANOSECOND:
            raise ibis_exceptions.UnsupportedOperationError(
                f"BigQuery does not allow binary operation {type(op).__name__} with "
                f"INTERVAL offset {unit}"
            )

        opname = type(op).__name__[len("Timestamp") :]
        funcname = f"TIMESTAMP_{opname.upper()}"
        return self.f.anon[funcname](left, right)

    visit_TimestampAdd = visit_TimestampSub = visit_TimestampAddSub

    def visit_DateAddSub(self, op, *, left, right):
        if not isinstance(right, sge.Interval):
            raise ibis_exceptions.OperationNotDefinedError(
                "BigQuery does not support non-literals on the right side of date add/subtract"
            )
        if not (unit := op.right.dtype.unit).is_date():
            raise ibis_exceptions.UnsupportedOperationError(
                f"BigQuery does not allow binary operation {type(op).__name__} with "
                f"INTERVAL offset {unit}"
            )
        opname = type(op).__name__[len("Date") :]
        funcname = f"DATE_{opname.upper()}"
        return self.f.anon[funcname](left, right)

    visit_DateAdd = visit_DateSub = visit_DateAddSub

    def visit_Covariance(self, op, *, left, right, how, where):
        if where is not None:
            left = self.if_(where, left, NULL)
            right = self.if_(where, right, NULL)

        if op.left.dtype.is_boolean():
            left = self.cast(left, dt.int64)

        if op.right.dtype.is_boolean():
            right = self.cast(right, dt.int64)

        how = op.how[:4].upper()
        assert how in ("POP", "SAMP"), 'how not in ("POP", "SAMP")'
        return self.agg[f"COVAR_{how}"](left, right, where=where)

    def visit_Correlation(self, op, *, left, right, how, where):
        if how == "sample":
            raise ValueError(f"Correlation with how={how!r} is not supported.")

        if where is not None:
            left = self.if_(where, left, NULL)
            right = self.if_(where, right, NULL)

        if op.left.dtype.is_boolean():
            left = self.cast(left, dt.int64)

        if op.right.dtype.is_boolean():
            right = self.cast(right, dt.int64)

        return self.agg.corr(left, right, where=where)

    def visit_TypeOf(self, op, *, arg):
        return self._pudf("typeof", arg)

    def visit_Xor(self, op, *, left, right):
        return sg.or_(sg.and_(left, sg.not_(right)), sg.and_(sg.not_(left), right))

    def visit_HashBytes(self, op, *, arg, how):
        if how not in ("md5", "sha1", "sha256", "sha512"):
            raise NotImplementedError(how)
        return self.f[how](arg)

    @staticmethod
    def _gen_valid_name(name: str) -> str:
        candidate = "_".join(map(str.strip, _NAME_REGEX.findall(name))) or "tmp"
        # column names cannot be longer than 300 characters
        #
        # https://cloud.google.com/bigquery/docs/schemas#column_names
        #
        # it's easy to rename columns, so raise an exception telling the user
        # to do so
        #
        # we could potentially relax this and support arbitrary-length columns
        # by compressing the information using hashing, but there's no reason
        # to solve that problem until someone encounters this error and cannot
        # rename their columns
        limit = 300
        if len(candidate) > limit:
            raise ibis_exceptions.IbisError(
                f"BigQuery does not allow column names longer than {limit:d} characters. "
                "Please rename your columns to have fewer characters."
            )
        return candidate

    def visit_CountStar(self, op, *, arg, where):
        if where is not None:
            return self.f.countif(where)
        return self.f.count(STAR)

    def visit_CountDistinctStar(self, op, *, where, arg):
        # Bigquery does not support count(distinct a,b,c) or count(distinct (a, b, c))
        # as expressions must be "groupable":
        # https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#group_by_grouping_item
        #
        # Instead, convert the entire expression to a string
        # SELECT COUNT(DISTINCT concat(to_json_string(a), to_json_string(b)))
        # This works with an array of datatypes which generates a unique string
        # https://cloud.google.com/bigquery/docs/reference/standard-sql/json_functions#json_encodings
        row = sge.Concat(
            expressions=[
                self.f.to_json_string(sg.column(x, quoted=self.quoted))
                for x in op.arg.schema.keys()
            ]
        )
        if where is not None:
            row = self.if_(where, row, NULL)
        return self.f.count(sge.Distinct(expressions=[row]))

    def visit_Degrees(self, op, *, arg):
        return self._pudf("degrees", arg)

    def visit_Radians(self, op, *, arg):
        return self._pudf("radians", arg)

    def visit_CountDistinct(self, op, *, arg, where):
        if where is not None:
            arg = self.if_(where, arg, NULL)
        return self.f.count(sge.Distinct(expressions=[arg]))

    def visit_RandomUUID(self, op, **kwargs):
        return self.f.generate_uuid()

    def visit_ExtractFile(self, op, *, arg):
        return self._pudf("cw_url_extract_file", arg)

    def visit_ExtractFragment(self, op, *, arg):
        return self._pudf("cw_url_extract_fragment", arg)

    def visit_ExtractPath(self, op, *, arg):
        return self._pudf("cw_url_extract_path", arg)

    def visit_ExtractProtocol(self, op, *, arg):
        return self._pudf("cw_url_extract_protocol", arg)

    def visit_ExtractQuery(self, op, *, arg, key):
        if key is not None:
            return self._pudf("cw_url_extract_parameter", arg, key)
        else:
            return self._pudf("cw_url_extract_query", arg)

    def _pudf(self, name, *args):
        name = sg.table(name, db="persistent_udfs", catalog="bigquery-public-data").sql(
            self.dialect
        )
        return self.f[name](*args)

    def visit_DropColumns(self, op, *, parent, columns_to_drop):
        quoted = self.quoted
        excludes = [sg.column(column, quoted=quoted) for column in columns_to_drop]
        star = sge.Star(**{"except": excludes})
        table = sg.to_identifier(parent.alias_or_name, quoted=quoted)
        column = sge.Column(this=star, table=table)
        return sg.select(column).from_(parent)

    def visit_TableUnnest(
        self, op, *, parent, column, offset: str | None, keep_empty: bool
    ):
        quoted = self.quoted

        column_alias = sg.to_identifier(
            util.gen_name("table_unnest_column"), quoted=quoted
        )

        selcols = []

        table = sg.to_identifier(parent.alias_or_name, quoted=quoted)

        opname = op.column.name
        overlaps_with_parent = opname in op.parent.schema
        computed_column = column_alias.as_(opname, quoted=quoted)

        # replace the existing column if the unnested column hasn't been
        # renamed
        #
        # e.g., table.unnest("x")
        if overlaps_with_parent:
            selcols.append(
                sge.Column(this=sge.Star(replace=[computed_column]), table=table)
            )
        else:
            selcols.append(sge.Column(this=STAR, table=table))
            selcols.append(computed_column)

        if offset is not None:
            offset = sg.to_identifier(offset, quoted=quoted)
            selcols.append(offset)

        unnest = sge.Unnest(
            expressions=[column],
            alias=sge.TableAlias(columns=[column_alias]),
            offset=offset,
        )
        return (
            sg.select(*selcols)
            .from_(parent)
            .join(unnest, join_type="CROSS" if not keep_empty else "LEFT")
        )

    def visit_TimestampBucket(self, op, *, arg, interval, offset):
        arg_dtype = op.arg.dtype
        if arg_dtype.timezone is not None:
            funcname = "timestamp"
        else:
            funcname = "datetime"

        func = self.f[f"{funcname}_bucket"]

        origin = sge.convert("1970-01-01")
        if offset is not None:
            origin = self.f.anon[f"{funcname}_add"](origin, offset)

        return func(arg, interval, origin)

    def _array_reduction(self, *, arg, reduction):
        name = sg.to_identifier(util.gen_name(f"bq_arr_{reduction}"))
        return (
            sg.select(self.f[reduction](name))
            .from_(self._unnest(arg, as_=name))
            .subquery()
        )

    def visit_ArrayMin(self, op, *, arg):
        return self._array_reduction(arg=arg, reduction="min")

    def visit_ArrayMax(self, op, *, arg):
        return self._array_reduction(arg=arg, reduction="max")

    def visit_ArraySum(self, op, *, arg):
        return self._array_reduction(arg=arg, reduction="sum")

    def visit_ArrayMean(self, op, *, arg):
        return self._array_reduction(arg=arg, reduction="avg")

    def visit_ArrayAny(self, op, *, arg):
        return self._array_reduction(arg=arg, reduction="logical_or")

    def visit_ArrayAll(self, op, *, arg):
        return self._array_reduction(arg=arg, reduction="logical_and")

    # Customized ops for bigframes

    def visit_InMemoryTable(self, op, *, name, schema, data):
        # Avoid creating temp tables for small data, which is how memtable is
        # used in BigQuery DataFrames. Inspired by:
        # https://github.com/ibis-project/ibis/blob/efa6fb72bf4c790450d00a926d7bd809dade5902/ibis/backends/druid/compiler.py#L95
        rows = data.to_pyarrow(schema=None).to_pylist()  # type: ignore
        quoted = self.quoted
        columns = [sg.column(col, quoted=quoted) for col in schema.names]
        array_expr = sge.DataType(
            this=sge.DataType.Type.STRUCT,
            expressions=[
                sge.ColumnDef(
                    this=sge.to_identifier(field, quoted=self.quoted),
                    kind=bq_datatypes.BigQueryType.from_ibis(type_),
                )
                for field, type_ in zip(schema.names, schema.types)
            ],
            nested=True,
        )
        array_values = [
            sge.Struct(
                expressions=tuple(
                    self.visit_Literal(None, value=value, dtype=type_)
                    for value, type_ in zip(row.values(), schema.types)
                )
            )
            for row in rows
        ]
        expr = sge.Unnest(
            expressions=[
                sge.DataType(
                    this=sge.DataType.Type.ARRAY,
                    expressions=[array_expr],
                    nested=True,
                    values=array_values,
                ),
            ],
            alias=sge.TableAlias(
                this=sg.to_identifier(name, quoted=quoted),
                columns=columns,
            ),
        )
        return sg.select(sge.Star()).from_(expr)

    def visit_ArrayAggregate(self, op, *, arg, order_by, where):
        if len(order_by) > 0:
            expr = sge.Order(
                this=arg,
                expressions=[
                    # Avoid adding NULLS FIRST / NULLS LAST in SQL, which is
                    # unsupported in ARRAY_AGG by reconstructing the node as
                    # plain SQL text.
                    f"({order_column.args['this'].sql(dialect='bigquery')}) {'DESC' if order_column.args.get('desc') else 'ASC'}"
                    for order_column in order_by
                ],
            )
        else:
            expr = arg
        return sge.IgnoreNulls(this=self.agg.array_agg(expr, where=where))

    def visit_StringAgg(self, op, *, arg, sep, order_by, where):
        if len(order_by) > 0:
            expr = sge.Order(
                this=arg,
                expressions=[
                    # Avoid adding NULLS FIRST / NULLS LAST in SQL, which is
                    # unsupported in ARRAY_AGG by reconstructing the node as
                    # plain SQL text.
                    f"({order_column.args['this'].sql(dialect='bigquery')}) {'DESC' if order_column.args.get('desc') else 'ASC'}"
                    for order_column in order_by
                ],
            )
        else:
            expr = arg
        return self.agg.string_agg(expr, sep, where=where)

    def visit_AIGenerate(self, op, **kwargs):
        return sge.func("AI.GENERATE", *self._compile_ai_args(**kwargs))

    def visit_AIGenerateBool(self, op, **kwargs):
        return sge.func("AI.GENERATE_BOOL", *self._compile_ai_args(**kwargs))

    def visit_AIGenerateInt(self, op, **kwargs):
        return sge.func("AI.GENERATE_INT", *self._compile_ai_args(**kwargs))

    def visit_AIGenerateDouble(self, op, **kwargs):
        return sge.func("AI.GENERATE_DOUBLE", *self._compile_ai_args(**kwargs))

    def visit_AIIf(self, op, **kwargs):
        return sge.func("AI.IF", *self._compile_ai_args(**kwargs))

    def visit_AIClassify(self, op, **kwargs):
        return sge.func("AI.CLASSIFY", *self._compile_ai_args(**kwargs))

    def visit_AIScore(self, op, **kwargs):
        return sge.func("AI.SCORE", *self._compile_ai_args(**kwargs))

    def _compile_ai_args(self, **kwargs):
        args = []

        for key, val in kwargs.items():
            if val is None:
                continue

            if key == "model_params":
                val = sge.JSON(this=val)

            args.append(sge.Kwarg(this=sge.Identifier(this=key), expression=val))

        return args

    def visit_FirstNonNullValue(self, op, *, arg):
        return sge.IgnoreNulls(this=sge.FirstValue(this=arg))

    def visit_LastNonNullValue(self, op, *, arg):
        return sge.IgnoreNulls(this=sge.LastValue(this=arg))

    def visit_ToJsonString(self, op, *, arg):
        return self.f.to_json_string(arg)

    def visit_Quantile(self, op, *, arg, quantile, where):
        return sge.PercentileCont(this=arg, expression=quantile)

    def visit_WindowFunction(self, op, *, how, func, start, end, group_by, order_by):
        # Patch for https://github.com/ibis-project/ibis/issues/9872

        if start is None:
            start = {}
        if end is None:
            end = {}

        start_value = start.get("value", "UNBOUNDED")
        start_side = start.get("side", "PRECEDING")
        end_value = end.get("value", "UNBOUNDED")
        end_side = end.get("side", "FOLLOWING")

        if getattr(start_value, "this", None) == "0":
            start_value = "CURRENT ROW"
            start_side = None

        if getattr(end_value, "this", None) == "0":
            end_value = "CURRENT ROW"
            end_side = None

        if how != "none":
            spec = sge.WindowSpec(
                kind=how.upper(),
                start=start_value,
                start_side=start_side,
                end=end_value,
                end_side=end_side,
                over="OVER",
            )
        else:
            spec = None

        # If unordered, unbound range window is implicit
        if (not order_by) and (not start) and (not end):
            spec = None

        order = sge.Order(expressions=order_by) if order_by else None

        return sge.Window(this=func, partition_by=group_by, order=order, spec=spec)


compiler = BigQueryCompiler()
