# Copyright 2023 Agentic.AI Corporation.
#
# 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.

"""Core resource interface."""

import abc
import functools
import json
from typing import (
  Dict, Generic, Iterable, List, Optional, Tuple, Type, TypeVar, Union)

from enact import contexts


JsonLeaf = Union[int, float, str, bool, None]
Json = Union[JsonLeaf, List['Json'], Dict[str, 'Json']]


PRIMITIVES = (int, float, str, bytes, bool, type(None))


Primitives = Union[
  int, float, str, bytes, bool,
  None]


FieldValue = Union[
  Primitives,
  'ResourceBase',
  Type['ResourceBase'],
  List['FieldValue'],
  Dict[str, 'FieldValue']]


ResourceDictValue = Union[
  Primitives,
  Type['ResourceBase'],
  List['ResourceDictValue'],
  Dict[str, 'ResourceDictValue'],
  'ResourceDict']


C = TypeVar('C', bound='ResourceBase')


class FrameworkError(Exception):
  """Superclass for framework related errors."""


class ResourceError(FrameworkError):
  """Superclass for resource related errors."""


class ImplementationMissing(FrameworkError):
  """Superclass for errors where an implementaton is missing."""


class FieldTypeError(FrameworkError):
  """Superclass for errors related to field types."""


@contexts.register
class _AcyclicContext(contexts.Context):
  """Helper to safeguard against cyclic data-structures."""

  def __init__(self, obj):
    """Initializes the context."""
    super().__init__()
    self.parent: Optional[_AcyclicContext] = None
    self.obj = obj

  def enter(self):
    """Check if the resource is already being committed."""
    self.parent = _AcyclicContext.get_current()
    parent = self.parent
    parents: List[_AcyclicContext] = []
    while parent is not None:
      parents.append(parent)
      if parent.obj is self.obj:
        raise FieldTypeError(
          f'Resources may not have cyclic graph structure. '
          f'Encountered cycle: '
          f'{" -> ".join(str(p.obj) for p in parents)}')
      parent = parent.parent


class ResourceBase:
  """Base class for resources.

  Not an abstract base class in order to avoid meta-class conflict.

  Resources have a unique type identifier. Each resource class is associated
  with a fixed list of named fields. These fields must have value semantics,
  that is, replacing a field by a copy should not change the meaning of the
  resource. This means, for example, that code using resources should not rely
  on aliasing assumptions (e.g., two resources sharing the same list instance).

  In particular, this also means that resources may not mutually reference each
  other.
  """

  @classmethod
  def type_descr(cls) -> Dict[str, Json]:
    """Returns a unique descriptor for the type."""
    return {'name': f'{cls.__module__}.{cls.__qualname__}'}

  @classmethod
  @functools.lru_cache
  def type_id(cls) -> str:
    """Returns a string descriptor of the type."""
    return json.dumps(cls.type_descr(), sort_keys=True)

  @classmethod
  @abc.abstractmethod
  def field_names(cls) -> Iterable[str]:
    """Returns the names of the fields of the resource."""
    raise NotImplementedError(f'{cls} does not implement field_names')

  @abc.abstractmethod
  def field_values(self) -> Iterable[FieldValue]:
    """Return a list of field values, aligned with field_names."""
    raise NotImplementedError(f'{type(self)} does not implement field_values')

  def field_items(self) -> Iterable[Tuple[str, FieldValue]]:
    """Iterate through the field names and values."""
    return zip(self.field_names(), self.field_values())

  @classmethod
  @abc.abstractmethod
  def from_fields(cls: Type[C],
                  field_dict: Dict[str, FieldValue]) -> C:
    """Constructs the resource from a field dictionary."""
    raise NotImplementedError()

  @staticmethod
  def _to_dict_value(value: FieldValue) -> ResourceDictValue:
    """Transforms a field value to a resource dict value."""
    with _AcyclicContext(value):
      if isinstance(value, ResourceBase):
        return value.to_resource_dict()
      if isinstance(value, PRIMITIVES):
        return value
      if isinstance(value, type) and issubclass(value, ResourceBase):
        return value
      if isinstance(value, List):
        return [ResourceBase._to_dict_value(x) for x in value]
      if isinstance(value, Dict):
        def _assert_str(maybe_str: str) -> str:
          if type(maybe_str) is not str:  # pylint: disable=unidiomatic-typecheck
            raise FieldTypeError(
              f'Expected string key, got {type(maybe_str)}')
          return maybe_str
        return {
          _assert_str(k): ResourceBase._to_dict_value(v)
          for k, v in value.items()}
      raise FieldTypeError(
        f'Encountered unsupported field type {type(value)}: {value}')

  @staticmethod
  def _from_dict_value(value: ResourceDictValue) -> FieldValue:
    """Transforms a resource dict value to a field value."""
    if isinstance(value, PRIMITIVES):
      return value
    if isinstance(value, type) and issubclass(value, ResourceBase):
      return value
    if isinstance(value, List):
      return [ResourceBase._from_dict_value(x) for x in value]
    if isinstance(value, ResourceDict):
      return value.type.from_resource_dict(value)
    if isinstance(value, Dict):
      def _assert_str(maybe_str: str) -> str:
        if type(maybe_str) is not str:  # pylint: disable=unidiomatic-typecheck
          raise FieldTypeError(
            f'Expected string key, got {type(maybe_str)}')
        return maybe_str
      return {
        _assert_str(k): ResourceBase._from_dict_value(v)
        for k, v in value.items()}
    raise FieldTypeError(
      f'Encountered unsupported resource '
      f'dict value type {type(value)}: {value}')

  def deepcopy_resource(self: C) -> C:
    """Create a deep-copy of the resource."""
    return self.from_resource_dict(self.to_resource_dict())

  def to_resource_dict(self: C) -> 'ResourceDict[C]':
    """Returns a ResourceDict dictionary representation."""
    result = ResourceDict(type(self))
    for field_name, value in self.field_items():
      result[field_name] = ResourceBase._to_dict_value(value)
    return result

  @classmethod
  def from_resource_dict(cls: Type[C], resorce_dict: 'ResourceDict') -> C:
    """Constructs the resource from a ResourceDict dictionary."""
    if not issubclass(resorce_dict.type, cls):
      raise ResourceError(
        f'Expected resource of type {cls}, got {resorce_dict.type}')
    field_dict = {
      k: ResourceBase._from_dict_value(v)
      for k, v in resorce_dict.items()}
    return cls.from_fields(field_dict)

  def set_from(self, other: 'ResourceBase'):
    """Sets the fields of this resource from another resource.

    Implementation of set_from is required to support replays of invokable
    resources that change their internal state during execution.

    Args:
      other: The resource to set fields from.
    """

    raise ImplementationMissing(
      f'Setting fields from another resource is not '
      f'supported by type {type(self)}.')


class ResourceDict(Generic[C], Dict[str, ResourceDictValue]):
  """A dictionary representing a resource with attached type info."""

  def __init__(self, resource_type: Type[C], *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.type = resource_type

  def to_resource(self) -> C:
    """Constructs the resource from the dictionary."""
    return self.type.from_resource_dict(self)


WrappedT = TypeVar('WrappedT')
WrapperT = TypeVar('WrapperT', bound='TypeWrapperBase')


class TypeWrapperBase(ResourceBase, Generic[WrappedT]):
  """Interface for resource classes that wrap python classes"""
  @classmethod
  @abc.abstractmethod
  def wrapped_type(cls) -> Type[WrappedT]:
    """Returns the type of the wrapped value."""
    raise NotImplementedError()

  @classmethod
  @abc.abstractmethod
  def wrap(cls: Type[WrapperT], value: WrappedT) -> WrapperT:
    """Wrap a value."""
    raise NotImplementedError()

  @abc.abstractmethod
  def unwrap(self) -> WrappedT:
    """Wrap a value."""
    raise NotImplementedError()

  @classmethod
  def is_immutable(cls) -> bool:
    """Whether the wrapped value is immutable."""
    return False

  @classmethod
  def set_wrapped_value(cls, target: WrappedT, src: WrappedT):
    """Set a wrapped value target to correspond to source."""
    raise ImplementationMissing(
      f'Please implement set_wrapped_value for TypeWrapper'
      f'{cls} to enable advanced features, e.g., replays.')

