# Copyright 2026 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.

from __future__ import annotations

from collections.abc import Callable
import functools
import os
from typing import Any
from typing import cast
from typing import Optional
from typing import overload
from typing import Protocol
from typing import TypeVar
import warnings

T = TypeVar("T")


class _FeatureDecorator(Protocol):
  """A feature decorator usable with or without a message argument.

  Preserves the decorated object's type so that subclasses and type
  checkers continue to see the real class/function rather than ``Any``.
  """

  # @decorator (bare, on a class or function)
  @overload
  def __call__(self, message_or_obj: T) -> T:
    ...

  # @decorator() or @decorator("message")
  @overload
  def __call__(self, message_or_obj: Optional[str] = ...) -> Callable[[T], T]:
    ...

  def __call__(self, message_or_obj: Any = None) -> Any:
    ...


def _is_truthy_env(var_name: str) -> bool:
  value = os.environ.get(var_name)
  if value is None:
    return False
  return value.strip().lower() in ("1", "true", "yes", "on")


def _make_feature_decorator(
    *,
    label: str,
    default_message: str,
    block_usage: bool = False,
    bypass_env_var: Optional[str] = None,
) -> _FeatureDecorator:
  def decorator_factory(message_or_obj: Any = None) -> Any:
    # Case 1: Used as @decorator without parentheses
    # message_or_obj is the decorated class/function
    if message_or_obj is not None and (
        isinstance(message_or_obj, type) or callable(message_or_obj)
    ):
      return _create_decorator(
          default_message, label, block_usage, bypass_env_var
      )(message_or_obj)

    # Case 2: Used as @decorator() with or without message
    # message_or_obj is either None or a string message
    message = (
        message_or_obj if isinstance(message_or_obj, str) else default_message
    )
    return _create_decorator(message, label, block_usage, bypass_env_var)

  return cast(_FeatureDecorator, decorator_factory)


def _create_decorator(
    message: str, label: str, block_usage: bool, bypass_env_var: Optional[str]
) -> Callable[[T], T]:
  def decorator(obj: T) -> T:
    obj_name = getattr(obj, "__name__", type(obj).__name__)
    msg = f"[{label.upper()}] {obj_name}: {message}"

    if isinstance(obj, type):  # decorating a class
      cls = cast(type[Any], obj)
      orig_init = cast(Any, cls).__init__

      @functools.wraps(orig_init)
      def new_init(self: Any, *args: Any, **kwargs: Any) -> Any:
        # Check if usage should be bypassed via environment variable at call time
        should_bypass = bypass_env_var is not None and _is_truthy_env(
            bypass_env_var
        )

        if should_bypass:
          # Bypass completely - no warning, no error
          pass
        elif block_usage:
          raise RuntimeError(msg)
        else:
          warnings.warn(msg, category=UserWarning, stacklevel=2)
        return orig_init(self, *args, **kwargs)

      cast(Any, cls).__init__ = new_init
      return cast(T, cls)

    elif callable(obj):  # decorating a function or method
      func = cast(Callable[..., Any], obj)

      @functools.wraps(func)
      def wrapper(*args: Any, **kwargs: Any) -> Any:
        # Check if usage should be bypassed via environment variable at call time
        should_bypass = bypass_env_var is not None and _is_truthy_env(
            bypass_env_var
        )

        if should_bypass:
          # Bypass completely - no warning, no error
          pass
        elif block_usage:
          raise RuntimeError(msg)
        else:
          warnings.warn(msg, category=UserWarning, stacklevel=2)
        return func(*args, **kwargs)

      return cast(T, wrapper)

    else:
      raise TypeError(
          f"@{label} can only be applied to classes or callable objects"
      )

  return decorator


working_in_progress = _make_feature_decorator(
    label="WIP",
    default_message=(
        "This feature is a work in progress and is not working completely. ADK"
        " users are not supposed to use it."
    ),
    block_usage=True,
    bypass_env_var="ADK_ALLOW_WIP_FEATURES",
)
"""Mark a class or function as a work in progress.

By default, decorated functions/classes will raise RuntimeError when used.
Set ADK_ALLOW_WIP_FEATURES=true environment variable to bypass this restriction.
ADK users are not supposed to set this environment variable.

Sample usage:

```
@working_in_progress("This feature is not ready for production use.")
def my_wip_function():
  pass
```
"""

experimental = _make_feature_decorator(
    label="EXPERIMENTAL",
    default_message=(
        "This feature is experimental and may change or be removed in future"
        " versions without notice. It may introduce breaking changes at any"
        " time."
    ),
    bypass_env_var="ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS",
)
"""Mark a class or a function as an experimental feature.

Sample usage:

```
# Use with default message
@experimental
class ExperimentalClass:
  pass

# Use with custom message
@experimental("This API may have breaking change in the future.")
class CustomExperimentalClass:
  pass

# Use with empty parentheses (same as default message)
@experimental()
def experimental_function():
  pass
```
"""
