# Copyright 2022 The jax3d Authors.
#
# 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.

"""Dataclass utils."""

import dataclasses
import enum
from typing import Any, Generic, Optional, Type, TypeVar, Union

from jax3d.utils.typing import Dataclass

_InT = TypeVar('_InT')
_OutT = TypeVar('_OutT')
_EnumT = TypeVar('_EnumT', bound=enum.Enum)


class DataclassField(Generic[_InT, _OutT]):
  """Abstract descriptor which perform data validation/conversion.

  This is somewhat similar to `traitlets`, but directly compatible with
  dataclasses.

  Example:

  ```
  class PositiveInteger(DataclassField[int, int]):

    def _default(self):
      return 123

    def _validate(self, value: int) -> int:
      if i < 0:
        raise ValueError('i should be positive')
      return i


  @dataclasses.dataclass
  class A:
    x: int = PositiveInteger()

  A(x=1)  # works
  A(x=-1)  # ValueError: i should be positive
  assert A().x == 123  # Default factory
  ```

  """

  def __init__(self, default: _InT = dataclasses.MISSING) -> None:
    """Constructor.

    Args:
      default: Default field value. If not passed, the value will be required.
    """
    # Attribute name and objtype refer to the object in which the descriptor
    # is applied. E.g. in the docstring example:
    # * _attribute_name = 'x'
    # * _objtype = A
    self._attribute_name: Optional[str] = None
    self._objtype: Optional[Type[Dataclass]] = None

    # Whether the descriptor is a required or not. If set to True,
    # calling `A()` will raise "Missing required kwargs 'x'" unless
    # `_default()` is overwritten.
    self._is_missing = default is dataclasses.MISSING

    # Default value (when `_is_missing is False`)
    # Should validate really called at import time ?
    self._default_value = None if self._is_missing else self._validate(default)

    # Whether `__get__` was called once. See `__get__` for details.
    self._first_getattr_call: bool = True

  def __set_name__(self, objtype: Type[Dataclass], name: str) -> None:
    self._objtype = objtype
    self._attribute_name = name

  def __get__(
      self,