# Copyright 2023 Google LLC
#
# Licensed 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.

"""DataFrame is a two dimensional data structure."""

from __future__ import annotations

import functools
from typing import Optional, Protocol, TYPE_CHECKING, Union

import bigframes_vendored.constants as constants

import bigframes.exceptions

if TYPE_CHECKING:
    from bigframes import Session
    from bigframes.core.blocks import Block
    from bigframes.dataframe import DataFrame
    from bigframes.series import Series


class HasSession(Protocol):
    @property
    def _session(self) -> Session:
        ...

    @property
    def _block(self) -> Block:
        ...


def requires_index(meth):
    @functools.wraps(meth)
    def guarded_meth(df: Union[DataFrame, Series], *args, **kwargs):
        df._throw_if_null_index(meth.__name__)
        return meth(df, *args, **kwargs)

    guarded_meth._validations_requires_index = True  # type: ignore
    return guarded_meth


def requires_ordering(suggestion: Optional[str] = None):
    def decorator(meth):
        @functools.wraps(meth)
        def guarded_meth(object: HasSession, *args, **kwargs):
            enforce_ordered(object, meth.__name__, suggestion)
            return meth(object, *args, **kwargs)

        guarded_meth._validations_requires_ordering = True  # type: ignore
        return guarded_meth

    return decorator


def enforce_ordered(
    object: HasSession, opname: str, suggestion: Optional[str] = None
) -> None:
    session = object._session
    if session._strictly_ordered or not object._block.expr.order_ambiguous:
        # No ambiguity for how to calculate ordering, so no error or warning
        return None
    if not session._allows_ambiguity:
        suggestion_substr = suggestion + " " if suggestion else ""
        raise bigframes.exceptions.OrderRequiredError(
            f"Op {opname} not supported when strict ordering is disabled. {suggestion_substr}{constants.FEEDBACK_LINK}"
        )
    if not object._block.explicitly_ordered:
        raise bigframes.exceptions.OrderRequiredError(
            f"Op {opname} requires an ordering. Use .sort_values or .sort_index to provide an ordering. {constants.FEEDBACK_LINK}"
        )
