# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from __future__ import annotations

from datetime import datetime, timezone

import pyarrow as pa
import pyarrow.compute as pc
import pytest
from datafusion import Accumulator, SessionContext, column, udaf
from datafusion.expr import Window


class Summarize(Accumulator):
    """Interface of a user-defined accumulation."""

    def __init__(self, initial_value: float = 0.0, as_scalar: bool = False):
        self._sum = initial_value
        self.as_scalar = as_scalar

    def state(self) -> list[pa.Scalar]:
        if self.as_scalar:
            return [pa.scalar(self._sum)]
        return [self._sum]

    def update(self, values: pa.Array) -> None:
        # Not nice since pyarrow scalars can't be summed yet.
        # This breaks on `None`
        self._sum = self._sum + pc.sum(values).as_py()

    def merge(self, states: list[pa.Array]) -> None:
        # Not nice since pyarrow scalars can't be summed yet.
        # This breaks on `None`
        self._sum = self._sum + pc.sum(states[0]).as_py()

    def evaluate(self) -> pa.Scalar:
        if self.as_scalar:
            return pa.scalar(self._sum)
        return self._sum


class NotSubclassOfAccumulator:
    pass


class MissingMethods(Accumulator):
    def __init__(self):
        self._sum = pa.scalar(0)

    def state(self) -> list[pa.Scalar]:
        return [self._sum]


class CollectTimestamps(Accumulator):
    def __init__(self, wrap_in_scalar: bool):
        self._values: list[datetime] = []
        self.wrap_in_scalar = wrap_in_scalar

    def state(self) -> list[pa.Scalar]:
        if self.wrap_in_scalar:
            return [pa.scalar(self._values, type=pa.list_(pa.timestamp("ns")))]
        return [pa.array(self._values, type=pa.timestamp("ns"))]

    def update(self, values: pa.Array) -> None:
        self._values.extend(values.to_pylist())

    def merge(self, states: list[pa.Array]) -> None:
        for state in states[0].to_pylist():
            if state is not None:
                self._values.extend(state)

    def evaluate(self) -> pa.Scalar:
        if self.wrap_in_scalar:
            return pa.scalar(self._values, type=pa.list_(pa.timestamp("ns")))
        return pa.array(self._values, type=pa.timestamp("ns"))


@pytest.fixture
def df(ctx):
    # create a RecordBatch and a new DataFrame from it
    batch = pa.RecordBatch.from_arrays(
        [pa.array([1, 2, 3]), pa.array([4, 4, 6])],
        names=["a", "b"],
    )
    return ctx.create_dataframe([[batch]], name="test_table")


def test_errors(df):
    with pytest.raises(TypeError):
        udaf(
            NotSubclassOfAccumulator,
            pa.float64(),
            pa.float64(),
            [pa.float64()],
            volatility="immutable",
        )

    msg = (
        "Can't instantiate abstract class MissingMethods (without an implementation "
        "for abstract methods 'evaluate', 'merge', 'update'|with abstract methods "
        "evaluate, merge, update)"
    )
    with pytest.raises(Exception, match=msg):
        accum = udaf(  # noqa: F841
            MissingMethods,
            pa.int64(),
            pa.int64(),
            [pa.int64()],
            volatility="immutable",
        )


def test_udaf_aggregate(df):
    summarize = udaf(
        Summarize,
        pa.float64(),
        pa.float64(),
        [pa.float64()],
        volatility="immutable",
    )

    df1 = df.aggregate([], [summarize(column("a"))])

    # execute and collect the first (and only) batch
    result = df1.collect()[0]

    assert result.column(0) == pa.array([1.0 + 2.0 + 3.0])

    df2 = df.aggregate([], [summarize(column("a"))])

    # Run a second time to ensure the state is properly reset
    result = df2.collect()[0]

    assert result.column(0) == pa.array([1.0 + 2.0 + 3.0])


def test_udaf_decorator_aggregate(df):
    @udaf(pa.float64(), pa.float64(), [pa.float64()], "immutable")
    def summarize():
        return Summarize()

    df1 = df.aggregate([], [summarize(column("a"))])

    # execute and collect the first (and only) batch
    result = df1.collect()[0]

    assert result.column(0) == pa.array([1.0 + 2.0 + 3.0])

    df2 = df.aggregate([], [summarize(column("a"))])

    # Run a second time to ensure the state is properly reset
    result = df2.collect()[0]

    assert result.column(0) == pa.array([1.0 + 2.0 + 3.0])


def test_udaf_decorator_keyword_arguments(df):
    @udaf(
        input_types=pa.float64(),
        return_type=pa.float64(),
        state_type=[pa.float64()],
        volatility="immutable",
    )
    def summarize():
        return Summarize()

    result = df.aggregate([], [summarize(column("a"))]).collect()[0]
    assert result.column(0) == pa.array([1.0 + 2.0 + 3.0])


def test_udaf_function_keyword_arguments(df):
    summarize = udaf(
        accum=Summarize,
        input_types=pa.float64(),
        return_type=pa.float64(),
        state_type=[pa.float64()],
        volatility="immutable",
    )
    result = df.aggregate([], [summarize(column("a"))]).collect()[0]
    assert result.column(0) == pa.array([1.0 + 2.0 + 3.0])


@pytest.mark.parametrize("as_scalar", [True, False])
def test_udaf_aggregate_with_arguments(df, as_scalar):
    bias = 10.0

    summarize = udaf(
        lambda: Summarize(initial_value=bias, as_scalar=as_scalar),
        pa.float64(),
        pa.float64(),
        [pa.float64()],
        volatility="immutable",
    )

    df1 = df.aggregate([], [summarize(column("a"))])

    # execute and collect the first (and only) batch
    result = df1.collect()[0]

    assert result.column(0) == pa.array([bias + 1.0 + 2.0 + 3.0])

    df2 = df.aggregate([], [summarize(column("a"))])

    # Run a second time to ensure the state is properly reset
    result = df2.collect()[0]

    assert result.column(0) == pa.array([bias + 1.0 + 2.0 + 3.0])


def test_udaf_decorator_aggregate_with_arguments(df):
    bias = 10.0

    @udaf(pa.float64(), pa.float64(), [pa.float64()], "immutable")
    def summarize():
        return Summarize(bias)

    df1 = df.aggregate([], [summarize(column("a"))])

    # execute and collect the first (and only) batch
    result = df1.collect()[0]

    assert result.column(0) == pa.array([bias + 1.0 + 2.0 + 3.0])

    df2 = df.aggregate([], [summarize(column("a"))])

    # Run a second time to ensure the state is properly reset
    result = df2.collect()[0]

    assert result.column(0) == pa.array([bias + 1.0 + 2.0 + 3.0])


def test_group_by(df):
    summarize = udaf(
        Summarize,
        pa.float64(),
        pa.float64(),
        [pa.float64()],
        volatility="immutable",
    )

    df = df.aggregate([column("b")], [summarize(column("a"))])

    batches = df.collect()

    arrays = [batch.column(1) for batch in batches]
    joined = pa.concat_arrays(arrays)
    assert joined == pa.array([1.0 + 2.0, 3.0])


def test_register_udaf(ctx, df) -> None:
    summarize = udaf(
        Summarize,
        pa.float64(),
        pa.float64(),
        [pa.float64()],
        volatility="immutable",
    )

    ctx.register_udaf(summarize)

    df_result = ctx.sql("select summarize(b) from test_table")

    assert df_result.collect()[0][0][0].as_py() == 14.0


@pytest.fixture
def distinct_ctx():
    ctx = SessionContext()
    ctx.register_udaf(
        udaf(
            Summarize,
            pa.float64(),
            pa.float64(),
            [pa.float64()],
            volatility="immutable",
        )
    )
    ctx.from_pydict({"v": [1.0, 1.0, 1.0, 5.0]}, name="t")
    return ctx


@pytest.mark.parametrize(
    "run",
    [
        pytest.param(
            lambda ctx, summarize: ctx.table("t").aggregate(
                [], [summarize(column("v")).distinct().build().alias("r")]
            ),
            id="distinct aggregate",
        ),
        pytest.param(
            lambda ctx, summarize: ctx.table("t").select(
                summarize(column("v")).over(Window()).distinct().build().alias("r")
            ),
            id="distinct window",
        ),
    ],
)
def test_udaf_distinct_raises(distinct_ctx, run):
    # The Python accumulator cannot deduplicate, so it would count every row.
    summarize = udaf(
        Summarize,
        pa.float64(),
        pa.float64(),
        [pa.float64()],
        volatility="immutable",
    )
    with pytest.raises(Exception, match="DISTINCT is not supported"):
        run(distinct_ctx, summarize).collect()


def test_udaf_distinct_rewritten_to_group_by_runs(distinct_ctx):
    # The optimizer groups by the distinct values, so the accumulator never
    # sees DISTINCT.
    result = distinct_ctx.sql("select summarize(distinct v) as r from t")
    assert result.collect_column("r").to_pylist() == [6.0]


@pytest.mark.parametrize("wrap_in_scalar", [True, False])
def test_udaf_list_timestamp_return(ctx, wrap_in_scalar) -> None:
    timestamps1 = [
        datetime(2024, 1, 1, tzinfo=timezone.utc),
        datetime(2024, 1, 2, tzinfo=timezone.utc),
    ]
    timestamps2 = [
        datetime(2024, 1, 3, tzinfo=timezone.utc),
        datetime(2024, 1, 4, tzinfo=timezone.utc),
    ]
    batch1 = pa.RecordBatch.from_arrays(
        [pa.array(timestamps1, type=pa.timestamp("ns"))],
        names=["ts"],
    )
    batch2 = pa.RecordBatch.from_arrays(
        [pa.array(timestamps2, type=pa.timestamp("ns"))],
        names=["ts"],
    )
    df = ctx.create_dataframe([[batch1], [batch2]], name="timestamp_table")

    list_type = pa.list_(
        pa.field("item", type=pa.timestamp("ns"), nullable=wrap_in_scalar)
    )

    collect = udaf(
        lambda: CollectTimestamps(wrap_in_scalar),
        pa.timestamp("ns"),
        list_type,
        [list_type],
        volatility="immutable",
    )

    result = df.aggregate([], [collect(column("ts"))]).collect()[0]

    # There is no guarantee about the ordering of the batches, so perform a sort
    # to get consistent results. Alternatively we could sort on evaluate().
    assert (
        result.column(0).values.sort()
        == pa.array(
            [[*timestamps1, *timestamps2]],
            type=list_type,
        ).values
    )