/*
 * Copyright (C) 2015 Square, Inc.
 *
 * 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.
 */
package com.squareup.leakcanary;

import androidx.annotation.NonNull;

import com.squareup.haha.perflib.ArrayInstance;
import com.squareup.haha.perflib.ClassInstance;
import com.squareup.haha.perflib.ClassObj;
import com.squareup.haha.perflib.Field;
import com.squareup.haha.perflib.HprofParser;
import com.squareup.haha.perflib.Instance;
import com.squareup.haha.perflib.RootObj;
import com.squareup.haha.perflib.RootType;
import com.squareup.haha.perflib.Snapshot;
import com.squareup.haha.perflib.Type;
import com.squareup.haha.perflib.io.HprofBuffer;
import com.squareup.haha.perflib.io.MemoryMappedFileBuffer;

import java.io.File;
import java.lang.reflect.Constructor;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;

import gnu.trove.THashMap;
import gnu.trove.TObjectProcedure;

import static android.os.Build.VERSION.SDK_INT;
import static android.os.Build.VERSION_CODES.N_MR1;
import static java.util.concurrent.TimeUnit.NANOSECONDS;

/**
 * Analyzes heap dumps generated by a {@link RefWatcher} to verify if suspected leaks are real.
 */
public final class HeapAnalyzer {

    private static final String ANONYMOUS_CLASS_NAME_PATTERN = "^.+\\$\\d+$";

    private final ExcludedRefs excludedRefs;
    private final AnalyzerProgressListener listener;
    private final List<Reachability.Inspector> reachabilityInspectors;

    /**
     * @deprecated Use {@link #HeapAnalyzer(ExcludedRefs, AnalyzerProgressListener, List)}.
     */
    @Deprecated
    public HeapAnalyzer(@NonNull ExcludedRefs excludedRefs) {
        this(excludedRefs, AnalyzerProgressListener.NONE,
                Collections.<Class<? extends Reachability.Inspector>>emptyList());
    }

    public HeapAnalyzer(@NonNull ExcludedRefs excludedRefs,
                        @NonNull AnalyzerProgressListener listener,
                        @NonNull List<Class<? extends Reachability.Inspector>> reachabilityInspectorClasses) {
        this.excludedRefs = excludedRefs;
        this.listener = listener;

        this.reachabilityInspectors = new ArrayList<>();
        for (Class<? extends Reachability.Inspector> reachabilityInspectorClass
                : reachabilityInspectorClasses) {
            try {
                Constructor<? extends Reachability.Inspector> defaultConstructor =
                        reachabilityInspectorClass.getDeclaredConstructor();
                reachabilityInspectors.add(defaultConstructor.newInstance());
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
    }

    public @NonNull
    List<TrackedReference> findTrackedReferences(@NonNull File heapDumpFile) {
        if (!heapDumpFile.exists()) {
            throw new IllegalArgumentException("File does not exist: " + heapDumpFile);
        }
        try {
            HprofBuffer buffer = new MemoryMappedFileBuffer(heapDumpFile);
            HprofParser parser = new HprofParser(buffer);
            Snapshot snapshot = parser.parse();
            deduplicateGcRoots(snapshot);

            ClassObj refClass = snapshot.findClass(KeyedWeakReference.class.getName());
            List<TrackedReference> references = new ArrayList<>();
            for (Instance weakRef : refClass.getInstancesList()) {
                List<ClassInstance.FieldValue> values = HahaHelper.classInstanceValues(weakRef);
                String key = HahaHelper.asString(HahaHelper.fieldValue(values, "key"));
                String name =
                        HahaHelper.hasField(values, "name") ? HahaHelper.asString(HahaHelper.fieldValue(values, "name")) : "(No name field)";
                Instance instance = HahaHelper.fieldValue(values, "referent");
                if (instance != null) {
                    String className = getClassName(instance);
                    List<LeakReference> fields = describeFields(instance);
                    references.add(new TrackedReference(key, name, className, fields));
                }
            }
            return references;
        } catch (Throwable e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * Calls {@link #checkForLeak(File, String, boolean)} with computeRetainedSize set to true.
     *
     * @deprecated Use {@link #checkForLeak(File, String, boolean)} instead.
     */
    @Deprecated
    public @NonNull
    AnalysisResult checkForLeak(@NonNull File heapDumpFile,
                                @NonNull String referenceKey) {
        return checkForLeak(heapDumpFile, referenceKey, true);
    }

    /**
     * Searches the heap dump for a {@link KeyedWeakReference} instance with the corresponding key,
     * and then computes the shortest strong reference path from that instance to the GC roots.
     */
    public @NonNull
    AnalysisResult checkForLeak(@NonNull File heapDumpFile,
                                @NonNull String referenceKey,
                                boolean computeRetainedSize) {
        long analysisStartNanoTime = System.nanoTime();

        if (!heapDumpFile.exists()) {
            Exception exception = new IllegalArgumentException("File does not exist: " + heapDumpFile);
            return AnalysisResult.failure(exception, since(analysisStartNanoTime));
        }

        try {
            listener.onProgressUpdate(AnalyzerProgressListener.Step.READING_HEAP_DUMP_FILE);
            HprofBuffer buffer = new MemoryMappedFileBuffer(heapDumpFile);
            HprofParser parser = new HprofParser(buffer);
            listener.onProgressUpdate(AnalyzerProgressListener.Step.PARSING_HEAP_DUMP);
            Snapshot snapshot = parser.parse();
            listener.onProgressUpdate(AnalyzerProgressListener.Step.DEDUPLICATING_GC_ROOTS);
            deduplicateGcRoots(snapshot);
            listener.onProgressUpdate(AnalyzerProgressListener.Step.FINDING_LEAKING_REF);
            Instance leakingRef = findLeakingReference(referenceKey, snapshot);

            // False alarm, weak reference was cleared in between key check and heap dump.
            if (leakingRef == null) {
                String className = leakingRef.getClassObj().getClassName();
                return AnalysisResult.noLeak(className, since(analysisStartNanoTime));
            }
            return findLeakTrace(analysisStartNanoTime, snapshot, leakingRef, computeRetainedSize);
        } catch (Throwable e) {
            return AnalysisResult.failure(e, since(analysisStartNanoTime));
        }
    }

    /**
     * 以前检查一遍 避免出现dokit内部的异常
     *
     * @param heapDumpFile
     * @param referenceKey
     * @return
     */
    public Instance preCheckForLeak(@NonNull File heapDumpFile,
                                    @NonNull String referenceKey) {

        if (!heapDumpFile.exists()) {
            return null;
        }

        Instance leakingRef;
        try {
            HprofBuffer buffer = new MemoryMappedFileBuffer(heapDumpFile);
            HprofParser parser = new HprofParser(buffer);
            Snapshot snapshot = parser.parse();
            deduplicateGcRoots(snapshot);
            leakingRef = findLeakingReference(referenceKey, snapshot);
        } catch (Exception e) {
            return null;
        }
        return leakingRef;
    }


    /**
     * Pruning duplicates reduces memory pressure from hprof bloat added in Marshmallow.
     */
    void deduplicateGcRoots(Snapshot snapshot) {
        // THashMap has a smaller memory footprint than HashMap.
        final THashMap<String, RootObj> uniqueRootMap = new THashMap<>();

        final Collection<RootObj> gcRoots = snapshot.getGCRoots();
        for (RootObj root : gcRoots) {
            String key = generateRootKey(root);
            if (!uniqueRootMap.containsKey(key)) {
                uniqueRootMap.put(key, root);
            }
        }

        // Repopulate snapshot with unique GC roots.
        gcRoots.clear();
        uniqueRootMap.forEach(new TObjectProcedure<String>() {
            @Override
            public boolean execute(String key) {
                return gcRoots.add(uniqueRootMap.get(key));
            }
        });
    }

    private String generateRootKey(RootObj root) {
        return String.format("%s@0x%08x", root.getRootType().getName(), root.getId());
    }

    private Instance findLeakingReference(String key, Snapshot snapshot) {
        ClassObj refClass = snapshot.findClass(KeyedWeakReference.class.getName());
        if (refClass == null) {
            throw new IllegalStateException(
                    "Could not find the " + KeyedWeakReference.class.getName() + " class in the heap dump.");
        }
        List<String> keysFound = new ArrayList<>();
        for (Instance instance : refClass.getInstancesList()) {
            List<ClassInstance.FieldValue> values = HahaHelper.classInstanceValues(instance);
            Object keyFieldValue = HahaHelper.fieldValue(values, "key");
            if (keyFieldValue == null) {
                keysFound.add(null);
                continue;
            }
            String keyCandidate = HahaHelper.asString(keyFieldValue);
            if (keyCandidate.equals(key)) {
                return HahaHelper.fieldValue(values, "referent");
            }
            keysFound.add(keyCandidate);
        }
        throw new IllegalStateException(
                "Could not find weak reference with key " + key + " in " + keysFound);
    }

    private AnalysisResult findLeakTrace(long analysisStartNanoTime, Snapshot snapshot,
                                         Instance leakingRef, boolean computeRetainedSize) {

        listener.onProgressUpdate(AnalyzerProgressListener.Step.FINDING_SHORTEST_PATH);
        ShortestPathFinder pathFinder = new ShortestPathFinder(excludedRefs);
        ShortestPathFinder.Result result = pathFinder.findPath(snapshot, leakingRef);

        String className = leakingRef.getClassObj().getClassName();

        // False alarm, no strong reference path to GC Roots.
        if (result.leakingNode == null) {
            return AnalysisResult.noLeak(className, since(analysisStartNanoTime));
        }

        listener.onProgressUpdate(AnalyzerProgressListener.Step.BUILDING_LEAK_TRACE);
        LeakTrace leakTrace = buildLeakTrace(result.leakingNode);

        long retainedSize;
        if (computeRetainedSize) {

            listener.onProgressUpdate(AnalyzerProgressListener.Step.COMPUTING_DOMINATORS);
            // Side effect: computes retained size.
            snapshot.computeDominators();

            Instance leakingInstance = result.leakingNode.instance;

            retainedSize = leakingInstance.getTotalRetainedSize();

            if (SDK_INT <= N_MR1) {
                listener.onProgressUpdate(AnalyzerProgressListener.Step.COMPUTING_BITMAP_SIZE);
                retainedSize += computeIgnoredBitmapRetainedSize(snapshot, leakingInstance);
            }
        } else {
            retainedSize = AnalysisResult.RETAINED_HEAP_SKIPPED;
        }

        return AnalysisResult.leakDetected(result.excludingKnownLeaks, className, leakTrace, retainedSize,
                since(analysisStartNanoTime));
    }

    /**
     * Bitmaps and bitmap byte arrays are sometimes held by native gc roots, so they aren't included
     * in the retained size because their root dominator is a native gc root.
     * To fix this, we check if the leaking instance is a dominator for each bitmap instance and then
     * add the bitmap size.
     * <p>
     * From experience, we've found that bitmap created in code (Bitmap.createBitmap()) are correctly
     * accounted for, however bitmaps set in layouts are not.
     */
    private long computeIgnoredBitmapRetainedSize(Snapshot snapshot, Instance leakingInstance) {
        long bitmapRetainedSize = 0;
        ClassObj bitmapClass = snapshot.findClass("android.graphics.Bitmap");

        for (Instance bitmapInstance : bitmapClass.getInstancesList()) {
            if (isIgnoredDominator(leakingInstance, bitmapInstance)) {
                ArrayInstance mBufferInstance = HahaHelper.fieldValue(HahaHelper.classInstanceValues(bitmapInstance), "mBuffer");
                // Native bitmaps have mBuffer set to null. We sadly can't account for them.
                if (mBufferInstance == null) {
                    continue;
                }
                long bufferSize = mBufferInstance.getTotalRetainedSize();
                long bitmapSize = bitmapInstance.getTotalRetainedSize();
                // Sometimes the size of the buffer isn't accounted for in the bitmap retained size. Since
                // the buffer is large, it's easy to detect by checking for bitmap size < buffer size.
                if (bitmapSize < bufferSize) {
                    bitmapSize += bufferSize;
                }
                bitmapRetainedSize += bitmapSize;
            }
        }
        return bitmapRetainedSize;
    }

    private boolean isIgnoredDominator(Instance dominator, Instance instance) {
        boolean foundNativeRoot = false;
        while (true) {
            Instance immediateDominator = instance.getImmediateDominator();
            if (immediateDominator instanceof RootObj
                    && ((RootObj) immediateDominator).getRootType() == RootType.UNKNOWN) {
                // Ignore native roots
                instance = instance.getNextInstanceToGcRoot();
                foundNativeRoot = true;
            } else {
                instance = immediateDominator;
            }
            if (instance == null) {
                return false;
            }
            if (instance == dominator) {
                return foundNativeRoot;
            }
        }
    }

    private LeakTrace buildLeakTrace(LeakNode leakingNode) {
        List<LeakTraceElement> elements = new ArrayList<>();
        // We iterate from the leak to the GC root
        LeakNode node = new LeakNode(null, null, leakingNode, null);
        while (node != null) {
            LeakTraceElement element = buildLeakElement(node);
            if (element != null) {
                elements.add(0, element);
            }
            node = node.parent;
        }

        List<Reachability> expectedReachability =
                computeExpectedReachability(elements);

        return new LeakTrace(elements, expectedReachability);
    }

    private List<Reachability> computeExpectedReachability(
            List<LeakTraceElement> elements) {
        int lastReachableElement = 0;
        int lastElementIndex = elements.size() - 1;
        int firstUnreachableElement = lastElementIndex;
        // No need to inspect the first and last element. We know the first should be reachable (gc
        // root) and the last should be unreachable (watched instance).
        elementLoop:
        for (int i = 1; i < lastElementIndex; i++) {
            LeakTraceElement element = elements.get(i);

            for (Reachability.Inspector reachabilityInspector : reachabilityInspectors) {
                Reachability reachability = reachabilityInspector.expectedReachability(element);
                if (reachability == Reachability.REACHABLE) {
                    lastReachableElement = i;
                    break;
                } else if (reachability == Reachability.UNREACHABLE) {
                    firstUnreachableElement = i;
                    break elementLoop;
                }
            }
        }

        List<Reachability> expectedReachability = new ArrayList<>();
        for (int i = 0; i < elements.size(); i++) {
            Reachability status;
            if (i <= lastReachableElement) {
                status = Reachability.REACHABLE;
            } else if (i >= firstUnreachableElement) {
                status = Reachability.UNREACHABLE;
            } else {
                status = Reachability.UNKNOWN;
            }
            expectedReachability.add(status);
        }
        return expectedReachability;
    }

    private LeakTraceElement buildLeakElement(LeakNode node) {
        if (node.parent == null) {
            // Ignore any root node.
            return null;
        }
        Instance holder = node.parent.instance;

        if (holder instanceof RootObj) {
            return null;
        }
        LeakTraceElement.Holder holderType;
        String className;
        String extra = null;
        List<LeakReference> leakReferences = describeFields(holder);

        className = getClassName(holder);

        List<String> classHierarchy = new ArrayList<>();
        classHierarchy.add(className);
        String rootClassName = Object.class.getName();
        if (holder instanceof ClassInstance) {
            ClassObj classObj = holder.getClassObj();
            while (!(classObj = classObj.getSuperClassObj()).getClassName().equals(rootClassName)) {
                classHierarchy.add(classObj.getClassName());
            }
        }

        if (holder instanceof ClassObj) {
            holderType = LeakTraceElement.Holder.CLASS;
        } else if (holder instanceof ArrayInstance) {
            holderType = LeakTraceElement.Holder.ARRAY;
        } else {
            ClassObj classObj = holder.getClassObj();
            if (HahaHelper.extendsThread(classObj)) {
                holderType = LeakTraceElement.Holder.THREAD;
                String threadName = HahaHelper.threadName(holder);
                extra = "(named '" + threadName + "')";
            } else if (className.matches(ANONYMOUS_CLASS_NAME_PATTERN)) {
                String parentClassName = classObj.getSuperClassObj().getClassName();
                if (rootClassName.equals(parentClassName)) {
                    holderType = LeakTraceElement.Holder.OBJECT;
                    try {
                        // This is an anonymous class implementing an interface. The API does not give access
                        // to the interfaces implemented by the class. We check if it's in the class path and
                        // use that instead.
                        Class<?> actualClass = Class.forName(classObj.getClassName());
                        Class<?>[] interfaces = actualClass.getInterfaces();
                        if (interfaces.length > 0) {
                            Class<?> implementedInterface = interfaces[0];
                            extra = "(anonymous implementation of " + implementedInterface.getName() + ")";
                        } else {
                            extra = "(anonymous subclass of java.lang.Object)";
                        }
                    } catch (ClassNotFoundException ignored) {
                    }
                } else {
                    holderType = LeakTraceElement.Holder.OBJECT;
                    // Makes it easier to figure out which anonymous class we're looking at.
                    extra = "(anonymous subclass of " + parentClassName + ")";
                }
            } else {
                holderType = LeakTraceElement.Holder.OBJECT;
            }
        }
        return new LeakTraceElement(node.leakReference, holderType, classHierarchy, extra,
                node.exclusion, leakReferences);
    }

    private List<LeakReference> describeFields(Instance instance) {
        List<LeakReference> leakReferences = new ArrayList<>();
        if (instance instanceof ClassObj) {
            ClassObj classObj = (ClassObj) instance;
            for (Map.Entry<Field, Object> entry : classObj.getStaticFieldValues().entrySet()) {
                String name = entry.getKey().getName();
                String stringValue = HahaHelper.valueAsString(entry.getValue());
                leakReferences.add(new LeakReference(LeakTraceElement.Type.STATIC_FIELD, name, stringValue));
            }
        } else if (instance instanceof ArrayInstance) {
            ArrayInstance arrayInstance = (ArrayInstance) instance;
            if (arrayInstance.getArrayType() == Type.OBJECT) {
                Object[] values = arrayInstance.getValues();
                for (int i = 0; i < values.length; i++) {
                    String name = Integer.toString(i);
                    String stringValue = HahaHelper.valueAsString(values[i]);
                    leakReferences.add(new LeakReference(LeakTraceElement.Type.ARRAY_ENTRY, name, stringValue));
                }
            }
        } else {
            ClassObj classObj = instance.getClassObj();
            for (Map.Entry<Field, Object> entry : classObj.getStaticFieldValues().entrySet()) {
                String name = entry.getKey().getName();
                String stringValue = HahaHelper.valueAsString(entry.getValue());
                leakReferences.add(new LeakReference(LeakTraceElement.Type.STATIC_FIELD, name, stringValue));
            }
            ClassInstance classInstance = (ClassInstance) instance;
            for (ClassInstance.FieldValue field : classInstance.getValues()) {
                String name = field.getField().getName();
                String stringValue = HahaHelper.valueAsString(field.getValue());
                leakReferences.add(new LeakReference(LeakTraceElement.Type.INSTANCE_FIELD, name, stringValue));
            }
        }
        return leakReferences;
    }

    private String getClassName(Instance instance) {
        String className;
        if (instance instanceof ClassObj) {
            ClassObj classObj = (ClassObj) instance;
            className = classObj.getClassName();
        } else if (instance instanceof ArrayInstance) {
            ArrayInstance arrayInstance = (ArrayInstance) instance;
            className = arrayInstance.getClassObj().getClassName();
        } else {
            ClassObj classObj = instance.getClassObj();
            className = classObj.getClassName();
        }
        return className;
    }

    private long since(long analysisStartNanoTime) {
        return NANOSECONDS.toMillis(System.nanoTime() - analysisStartNanoTime);
    }
}
