/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.
 */
package org.apache.lucene.util;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.lang.reflect.Array;
import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.file.Path;
import java.security.AccessController;
import java.security.PrivilegedAction;
import java.util.AbstractList;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.function.Predicate;
import java.util.function.ToLongFunction;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

/** Crawls object graph to collect RAM usage for testing */
public final class RamUsageTester {

  /** An accumulator of object references. This class allows for customizing RAM usage estimation. */
  public static class Accumulator {

    /** Accumulate transitive references for the provided fields of the given
     *  object into <code>queue</code> and return the shallow size of this object. */
    public long accumulateObject(Object o, long shallowSize, Map<Field, Object> fieldValues, Collection<Object> queue) {
      queue.addAll(fieldValues.values());
      return shallowSize;
    }

    /** Accumulate transitive references for the provided values of the given
     *  array into <code>queue</code> and return the shallow size of this array. */
    public long accumulateArray(Object array, long shallowSize, List<Object> values, Collection<Object> queue) {
      queue.addAll(values);
      return shallowSize;
    }

  }

  /**
   * Estimates the RAM usage by the given object. It will
   * walk the object tree and sum up all referenced objects.
   *
   * <p><b>Resource Usage:</b> This method internally uses a set of
   * every object seen during traversals so it does allocate memory
   * (it isn't side-effect free). After the method exits, this memory
   * should be GCed.</p>
   */
  public static long sizeOf(Object obj, Accumulator accumulator) {
    return measureObjectSize(obj, accumulator);
  }

  /** Same as calling <code>sizeOf(obj, DEFAULT_FILTER)</code>. */
  public static long sizeOf(Object obj) {
    return sizeOf(obj, new Accumulator());
  }

  /**
   * Return a human-readable size of a given object.
   * @see #sizeOf(Object)
   * @see RamUsageEstimator#humanReadableUnits(long)
   */
  public static String humanSizeOf(Object object) {
    return RamUsageEstimator.humanReadableUnits(sizeOf(object));
  }

  /*
   * Non-recursive version of object descend. This consumes more memory than recursive in-depth
   * traversal but prevents stack overflows on long chains of objects
   * or complex graphs (a max. recursion depth on my machine was ~5000 objects linked in a chain
   * so not too much).
   */
  private static long measureObjectSize(Object root, Accumulator accumulator) {
    // Objects seen so far.
    final Set<Object> seen = Collections.newSetFromMap(new IdentityHashMap<Object, Boolean>());
    // Class cache with reference Field and precalculated shallow size. 
    final IdentityHashMap<Class<?>, ClassCache> classCache = new IdentityHashMap<>();
    // Stack of objects pending traversal. Recursion caused stack overflows. 
    final ArrayList<Object> stack = new ArrayList<>();
    stack.add(root);

    long totalSize = 0;
    while (!stack.isEmpty()) {
      final Object ob = stack.remove(stack.size() - 1);

      if (ob == null || seen.contains(ob)) {
        continue;
      }
      seen.add(ob);

      final long obSize;
      final Class<?> obClazz = ob.getClass();
      assert obClazz != null : "jvm bug detected (Object.getClass() == null). please report this to your vendor";
      if (obClazz.isArray()) {
        obSize = handleArray(accumulator, stack, ob, obClazz);
      } else {
        obSize = handleOther(accumulator, classCache, stack, ob, obClazz);
      }

      totalSize += obSize;
      // Dump size of each object for comparisons across JVMs and flags.
      // System.out.println("  += " + obClazz + " | " + obSize);
    }

    // Help the GC (?).
    seen.clear();
    stack.clear();
    classCache.clear();

    return totalSize;
  }

  private static long handleOther(Accumulator accumulator, IdentityHashMap<Class<?>, ClassCache> classCache, ArrayList<Object> stack, Object ob, Class<?> obClazz) {
    /*
     * Consider an object. Push any references it has to the processing stack
     * and accumulate this object's shallow size.
     */
    try {
      if (Constants.JRE_IS_MINIMUM_JAVA9) {
        long alignedShallowInstanceSize = RamUsageEstimator.shallowSizeOf(ob);

        Predicate<Class<?>> isJavaModule = (clazz) -> {
          return clazz.getName().startsWith("java.");
        };

        // Java 9: Best guess for some known types, as we cannot precisely look into runtime classes:
        final ToLongFunction<Object> func = SIMPLE_TYPES.get(obClazz);
        if (func != null) { // some simple type like String where the size is easy to get from public properties
          return accumulator.accumulateObject(ob, alignedShallowInstanceSize + func.applyAsLong(ob),
              Collections.emptyMap(), stack);
        } else if (ob instanceof Enum) {
          return alignedShallowInstanceSize;
        } else if (ob instanceof ByteBuffer) {
          // Approximate ByteBuffers with their underlying storage (ignores field overhead).
          return byteArraySize(((ByteBuffer) ob).capacity());
        }  else if (isJavaModule.test(obClazz) && ob instanceof Map) {
          final List<Object> values = ((Map<?,?>) ob).entrySet().stream()
              .flatMap(e -> Stream.of(e.getKey(), e.getValue()))
              .collect(Collectors.toList());
          return accumulator.accumulateArray(ob, alignedShallowInstanceSize + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER, values, stack)
              + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER;
        } else if (isJavaModule.test(obClazz) && ob instanceof Iterable) {
          final List<Object> values = StreamSupport.stream(((Iterable<?>) ob).spliterator(), false)
            .collect(Collectors.toList());
          return accumulator.accumulateArray(ob, alignedShallowInstanceSize + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER, values, stack)
              + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER;
        } else {
          // Fallback to reflective access.
        }
      }

      ClassCache cachedInfo = classCache.get(obClazz);
      if (cachedInfo == null) {
        classCache.put(obClazz, cachedInfo = createCacheEntry(obClazz));
      }

      final Map<Field, Object> fieldValues = new HashMap<>();
      for (Field f : cachedInfo.referenceFields) {
        fieldValues.put(f, f.get(ob));
      }
      return accumulator.accumulateObject(ob, cachedInfo.alignedShallowInstanceSize, fieldValues, stack);
    } catch (IllegalAccessException e) {
      // this should never happen as we enabled setAccessible().
      throw new RuntimeException("Reflective field access failed?", e);
    }
  }

  private static long handleArray(Accumulator accumulator, ArrayList<Object> stack, Object ob, Class<?> obClazz) {
    /*
     * Consider an array, possibly of primitive types. Push any of its references to
     * the processing stack and accumulate this array's shallow size.
     */
    final long shallowSize = RamUsageEstimator.shallowSizeOf(ob);
    final int len = Array.getLength(ob);
    final List<Object> values;
    Class<?> componentClazz = obClazz.getComponentType();
    if (componentClazz.isPrimitive()) {
      values = Collections.emptyList();
    } else {
      values = new AbstractList<Object>() {

        @Override
        public Object get(int index) {
          return Array.get(ob, index);
        }

        @Override
        public int size() {
          return len;
        }

      };
    }
    return accumulator.accumulateArray(ob, shallowSize, values, stack);
  }

  /**
   * This map contains a function to calculate sizes of some "simple types" like String just from their public properties.
   * This is needed for Java 9, which does not allow to look into runtime class fields.
   */
  @SuppressWarnings("serial")
  private static final Map<Class<?>, ToLongFunction<Object>> SIMPLE_TYPES = Collections.unmodifiableMap(new IdentityHashMap<Class<?>, ToLongFunction<Object>>() {
    { init(); }
    
    @SuppressForbidden(reason = "We measure some forbidden classes")
    private void init() {
      // String types:
      a(String.class, v -> charArraySize(v.length())); // may not be correct with Java 9's compact strings!
      a(StringBuilder.class, v -> charArraySize(v.capacity()));
      a(StringBuffer.class, v -> charArraySize(v.capacity()));
      // Types with large buffers:
      a(ByteArrayOutputStream.class, v -> byteArraySize(v.size()));
      // For File and Path, we just take the length of String representation as approximation:
      a(File.class, v -> charArraySize(v.toString().length()));
      a(Path.class, v -> charArraySize(v.toString().length()));
      a(ByteOrder.class, v -> 0); // Instances of ByteOrder are constants
    }
    
    @SuppressWarnings("unchecked")
    private <T> void a(Class<T> clazz, ToLongFunction<T> func) {
      put(clazz, (ToLongFunction<Object>) func);
    }
    
    private long charArraySize(int len) {
      return RamUsageEstimator.alignObjectSize((long)RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + (long)Character.BYTES * len);
    }
  });
  
  /**
   * Cached information about a given class.   
   */
  private static final class ClassCache {
    public final long alignedShallowInstanceSize;
    public final Field[] referenceFields;

    public ClassCache(long alignedShallowInstanceSize, Field[] referenceFields) {
      this.alignedShallowInstanceSize = alignedShallowInstanceSize;
      this.referenceFields = referenceFields;
    }    
  }
  
  /**
   * Create a cached information about shallow size and reference fields for 
   * a given class.
   */
  @SuppressForbidden(reason = "We need to access private fields of measured objects.")
  private static ClassCache createCacheEntry(final Class<?> clazz) {
    return AccessController.doPrivileged((PrivilegedAction<ClassCache>) () -> {
      ClassCache cachedInfo;
      long shallowInstanceSize = RamUsageEstimator.NUM_BYTES_OBJECT_HEADER;
      final ArrayList<Field> referenceFields = new ArrayList<>(32);
      for (Class<?> c = clazz; c != null; c = c.getSuperclass()) {
        if (c == Class.class) {
          // prevent inspection of Class' fields, throws SecurityException in Java 9!
          continue; 
        }
        final Field[] fields = c.getDeclaredFields();
        for (final Field f : fields) {
          if (!Modifier.isStatic(f.getModifiers())) {
            shallowInstanceSize = RamUsageEstimator.adjustForField(shallowInstanceSize, f);
  
            if (!f.getType().isPrimitive()) {
              try {
                f.setAccessible(true);
                referenceFields.add(f);
              } catch (RuntimeException re) {
                throw new RuntimeException(String.format(Locale.ROOT,
                    "Can't access field '%s' of class '%s' for RAM estimation.",
                    f.getName(),
                    clazz.getName()), re);
              }
            }
          }
        }
      }
  
      cachedInfo = new ClassCache(
          RamUsageEstimator.alignObjectSize(shallowInstanceSize), 
          referenceFields.toArray(new Field[referenceFields.size()]));
      return cachedInfo;
    });
  }

  private static long byteArraySize(int len) {
    return RamUsageEstimator.alignObjectSize((long) RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + len);
  }
}
