package arez.component.internal;

import arez.ActionFlags;
import arez.Arez;
import arez.ArezContext;
import arez.Component;
import arez.ComputableValue;
import arez.Disposable;
import arez.EqualityComparator;
import arez.ObjectsEqualsComparator;
import arez.SafeFunction;
import arez.Task;
import arez.component.DisposeNotifier;
import grim.annotations.OmitSymbol;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.Stack;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import org.intellij.lang.annotations.MagicConstant;
import static org.realityforge.braincheck.Guards.*;

/**
 * The class responsible for caching ComputableValue instances for different input parameters.
 */
public final class MemoizeCache<T>
  implements Disposable
{
  /**
   * Functional interface for calculating memoizable value.
   *
   * @param <T> The type of the returned value.
   */
  @FunctionalInterface
  public interface Function<T>
  {
    /**
     * Return calculated memoizable value.
     *
     * @param args the functions arguments.
     * @return the value generated by function.
     */
    T call( @Nonnull final Object... args );
  }

  /**
   * Reference to the system to which this node belongs.
   */
  @OmitSymbol( unless = "arez.enable_zones" )
  @Nullable
  private final ArezContext _context;
  /**
   * A human consumable prefix for computable values. It should be non-null if {@link Arez#areNamesEnabled()} returns
   * true and <code>null</code> otherwise.
   */
  @Nullable
  @OmitSymbol( unless = "arez.enable_names" )
  private final String _name;
  /**
   * The component that this memoize cache is contained within.
   * This should only be set if {@link Arez#areNativeComponentsEnabled()} is true but can be null even if this is true.
   */
  @OmitSymbol( unless = "arez.enable_native_components" )
  @Nullable
  private final Component _component;
  /**
   * The function memoized.
   */
  @Nonnull
  private final Function<T> _function;
  /**
   * The cache of all the ComputableValue created for each unique combination of parameters.
   */
  private final Map<Object, Object> _cache = new HashMap<>();
  /**
   * The number of arguments passed to memoized function.
   */
  private final int _argCount;
  /**
   * The flags passed to the created ComputableValues.
   */
  @MagicConstant( flagsFromClass = ComputableValue.Flags.class )
  private final int _flags;
  /**
   * Strategy used to compare old and new computed values.
   */
  @Nonnull
  private final EqualityComparator _equalityComparator;
  /**
   * The index of the next ComputableValue created.
   * This is only used when creating unique names for ComputableValues.
   */
  private int _nextIndex;
  /**
   * Flag indicating that the cache is currently being disposed.
   */
  private boolean _disposed;

  /**
   * Create the Memoize method cache.
   *
   * @param context   the context in which to create ComputableValue instances.
   * @param component the associated native component if any. This should only be set if {@link Arez#areNativeComponentsEnabled()} returns true.
   * @param name      a human consumable prefix for computable values.
   * @param function  the memoized function.
   * @param argCount  the number of arguments expected to be passed to memoized function.
   */
  public MemoizeCache( @Nullable final ArezContext context,
                       @Nullable final Component component,
                       @Nullable final String name,
                       @Nonnull final Function<T> function,
                       final int argCount )
  {
    this( context, component, name, function, argCount, 0, new ObjectsEqualsComparator() );
  }

  /**
   * Create the Memoize method cache.
   *
   * @param context   the context in which to create ComputableValue instances.
   * @param component the associated native component if any. This should only be set if {@link Arez#areNativeComponentsEnabled()} returns true.
   * @param name      a human consumable prefix for computable values.
   * @param function  the memoized function.
   * @param argCount  the number of arguments expected to be passed to memoized function.
   * @param flags     the flags that are used when creating ComputableValue instances. The only flags supported are flags defined in {@link ComputableValue.Flags} except for {@link ComputableValue.Flags#KEEPALIVE}, {@link ComputableValue.Flags#RUN_NOW} and {@link ComputableValue.Flags#RUN_LATER}.
   */
  public MemoizeCache( @Nullable final ArezContext context,
                       @Nullable final Component component,
                       @Nullable final String name,
                       @Nonnull final Function<T> function,
                       final int argCount,
                       @MagicConstant( flagsFromClass = ComputableValue.Flags.class ) final int flags )
  {
    this( context, component, name, function, argCount, flags, new ObjectsEqualsComparator() );
  }

  /**
   * Create the Memoize method cache.
   *
   * @param context            the context in which to create ComputableValue instances.
   * @param component          the associated native component if any. This should only be set if {@link Arez#areNativeComponentsEnabled()} returns true.
   * @param name               a human consumable prefix for computable values.
   * @param function           the memoized function.
   * @param argCount           the number of arguments expected to be passed to memoized function.
   * @param flags              the flags that are used when creating ComputableValue instances. The only flags supported are flags defined in {@link ComputableValue.Flags} except for {@link ComputableValue.Flags#KEEPALIVE}, {@link ComputableValue.Flags#RUN_NOW} and {@link ComputableValue.Flags#RUN_LATER}.
   * @param equalityComparator strategy used to compare old and new computed values.
   */
  public MemoizeCache( @Nullable final ArezContext context,
                       @Nullable final Component component,
                       @Nullable final String name,
                       @Nonnull final Function<T> function,
                       final int argCount,
                       @MagicConstant( flagsFromClass = ComputableValue.Flags.class ) final int flags,
                       @Nonnull final EqualityComparator equalityComparator )
  {
    if ( Arez.shouldCheckApiInvariants() )
    {
      apiInvariant( () -> Arez.areZonesEnabled() || null == context,
                    () -> "Arez-174: MemoizeCache passed a context but Arez.areZonesEnabled() is false" );
      apiInvariant( () -> Arez.areNamesEnabled() || null == name,
                    () -> "Arez-0159: MemoizeCache passed a name '" + name + "' but Arez.areNamesEnabled() is false" );
      apiInvariant( () -> argCount > 0,
                    () -> "Arez-0160: MemoizeCache constructed with invalid argCount: " + argCount +
                          ". Expected positive value." );
      final int mask = ComputableValue.Flags.PRIORITY_HIGHEST |
                       ComputableValue.Flags.PRIORITY_HIGH |
                       ComputableValue.Flags.PRIORITY_NORMAL |
                       ComputableValue.Flags.PRIORITY_LOW |
                       ComputableValue.Flags.PRIORITY_LOWEST |
                       ComputableValue.Flags.NO_REPORT_RESULT |
                       ComputableValue.Flags.AREZ_DEPENDENCIES |
                       ComputableValue.Flags.AREZ_OR_NO_DEPENDENCIES |
                       ComputableValue.Flags.AREZ_OR_EXTERNAL_DEPENDENCIES |
                       ComputableValue.Flags.OBSERVE_LOWER_PRIORITY_DEPENDENCIES |
                       ComputableValue.Flags.READ_OUTSIDE_TRANSACTION;

      apiInvariant( () -> ( ~mask & flags ) == 0,
                    () -> "Arez-0211: MemoizeCache passed unsupported flags. Unsupported bits: " + ( ~mask & flags ) );
    }
    _context = Arez.areZonesEnabled() ? Objects.requireNonNull( context ) : null;
    _component = Arez.areNativeComponentsEnabled() ? component : null;
    _name = Arez.areNamesEnabled() ? Objects.requireNonNull( name ) : null;
    _function = Objects.requireNonNull( function );
    _argCount = argCount;
    _flags = flags;
    _equalityComparator = Objects.requireNonNull( equalityComparator );
  }

  /**
   * Return the result of the memoized function, calculating if necessary.
   *
   * @param args the arguments passed to the memoized function.
   * @return the result of the memoized function.
   */
  public T get( @Nonnull final Object... args )
  {
    if ( Arez.shouldCheckApiInvariants() )
    {
      apiInvariant( this::isNotDisposed,
                    () -> "Arez-0161: MemoizeCache named '" + _name + "' had get() invoked when disposed." );
    }
    return getComputableValue( args ).get();
  }

  @Override
  public boolean isDisposed()
  {
    return _disposed;
  }

  @Override
  public void dispose()
  {
    if ( !_disposed )
    {
      _disposed = true;
      getContext().safeAction( Arez.areNamesEnabled() ? _name : null, () -> {
        disposeMap( _cache, _argCount );
        _cache.clear();
      }, ActionFlags.NO_VERIFY_ACTION_REQUIRED );
    }
  }

  @Nonnull
  private ArezContext getContext()
  {
    return Arez.areZonesEnabled() ? Objects.requireNonNull( _context ) : Arez.context();
  }

  /**
   * Traverse to leaf map elements and dispose all contained ComputableValue instances.
   */
  @SuppressWarnings( "unchecked" )
  private void disposeMap( @Nonnull final Map<Object, Object> map, final int depth )
  {
    if ( 1 == depth )
    {
      for ( final Map.Entry<Object, Object> entry : map.entrySet() )
      {
        final ComputableValue<?> computableValue = (ComputableValue<?>) entry.getValue();
        computableValue.dispose();
      }
    }
    else
    {
      for ( final Map.Entry<Object, Object> entry : map.entrySet() )
      {
        disposeMap( (Map<Object, Object>) entry.getValue(), depth - 1 );
      }
    }
  }

  /**
   * Retrieve the computable value for specified parameters, creating it if necessary.
   *
   * @param args the arguments passed to the memoized function.
   * @return the computable value instance for the specified args.
   */
  @SuppressWarnings( "unchecked" )
  @Nonnull
  public ComputableValue<T> getComputableValue( @Nonnull final Object... args )
  {
    if ( Arez.shouldCheckApiInvariants() )
    {
      apiInvariant( () -> args.length == _argCount,
                    () -> "Arez-0162: MemoizeCache.getComputableValue called with " + args.length +
                          " arguments but expected " + _argCount + " arguments." );
    }
    Map<Object, Object> map = _cache;
    final int size = args.length - 1;
    for ( int i = 0; i < size; i++ )
    {
      map = (Map<Object, Object>) map.computeIfAbsent( args[ i ], v -> new HashMap<>() );
    }
    ComputableValue<T> computableValue =
      (ComputableValue<T>) map.computeIfAbsent( args[ size ], v -> createComputableValue( args ) );
    if ( Disposable.isDisposed( computableValue ) )
    {
      computableValue = createComputableValue( args );
      map.put( args[ size ], computableValue );
    }
    return computableValue;
  }

  /**
   * Create computable value for specified parameters.
   *
   * @param args the arguments passed to the memoized function.
   */
  @Nonnull
  private ComputableValue<T> createComputableValue( @Nonnull final Object... args )
  {
    final Component component = Arez.areNativeComponentsEnabled() ? _component : null;
    final int id = _nextIndex++;
    final String name = Arez.areNamesEnabled() ? _name + "." + id : null;
    final SafeFunction<T> function = () -> {
      Arez.context().registerHook( "$MC$", null, () -> disposeComputableValue( args ) );
      return _function.call( args );
    };
    final ComputableValue<T> computable = getContext().computable( component, name, function, _flags, _equalityComparator );
    for ( final Object arg : args )
    {
      if ( arg instanceof DisposeNotifier )
      {
        DisposeNotifier.asDisposeNotifier( arg ).addOnDisposeListener( computable, computable::dispose, false );
      }
    }
    return computable;
  }

  /**
   * Method invoked to dispose memoized value.
   * This is called from deactivate hook so there should always by a cached value present
   * and thus we never check for missing elements in chain.
   *
   * @param args the arguments originally passed to the memoized function.
   */
  @SuppressWarnings( "unchecked" )
  void disposeComputableValue( @Nonnull final Object... args )
  {
    if ( Arez.shouldCheckInvariants() )
    {
      invariant( () -> args.length == _argCount,
                 () -> "Arez-0163: MemoizeCache.disposeComputableValue called with " + args.length +
                       " argument(s) but expected " + _argCount + " argument(s)." );
    }
    if ( _disposed )
    {
      return;
    }
    final Stack<Map<Object, ?>> stack = new Stack<>();
    stack.push( _cache );
    final int size = args.length - 1;
    for ( int i = 0; i < size; i++ )
    {
      stack.push( (Map<Object, ?>) stack.peek().get( args[ i ] ) );
    }
    final ComputableValue<T> computableValue = (ComputableValue<T>) stack.peek().remove( args[ size ] );
    if ( Arez.shouldCheckInvariants() )
    {
      invariant( () -> null != computableValue,
                 () -> "Arez-0193: MemoizeCache.disposeComputableValue called with args " + Arrays.asList( args ) +
                       " but unable to locate corresponding ComputableValue." );
    }
    assert null != computableValue;
    getContext().task( Arez.areNamesEnabled() ? computableValue.getName() + ".dispose" : null,
                       computableValue::dispose,
                       Task.Flags.PRIORITY_HIGHEST | Task.Flags.DISPOSE_ON_COMPLETE | Task.Flags.NO_WRAP_TASK );
    for ( final Object arg : args )
    {
      if ( arg instanceof DisposeNotifier )
      {
        DisposeNotifier.asDisposeNotifier( arg ).removeOnDisposeListener( computableValue, false );
      }
    }
    while ( stack.size() > 1 )
    {
      final Map<Object, ?> map = stack.pop();
      if ( map.isEmpty() )
      {
        stack.peek().remove( args[ stack.size() - 1 ] );
      }
      else
      {
        return;
      }
    }
  }

  @OmitSymbol
  Map<Object, Object> getCache()
  {
    return _cache;
  }

  @OmitSymbol
  int getNextIndex()
  {
    return _nextIndex;
  }

  @OmitSymbol
  int getFlags()
  {
    return _flags;
  }

  @OmitSymbol
  @Nonnull
  EqualityComparator getEqualityComparator()
  {
    return _equalityComparator;
  }
}