/*
 *  ******************************************************************************
 *  *
 *  *
 *  * This program and the accompanying materials are made available under the
 *  * terms of the Apache License, Version 2.0 which is available at
 *  * https://www.apache.org/licenses/LICENSE-2.0.
 *  *
 *  *  See the NOTICE file distributed with this work for additional
 *  *  information regarding copyright ownership.
 *  * 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.
 *  *
 *  * SPDX-License-Identifier: Apache-2.0
 *  *****************************************************************************
 */

package org.nd4j.linalg.api.memory.deallocation;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.Getter;
import lombok.NonNull;
import lombok.Setter;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import lombok.val;

import org.apache.commons.lang3.RandomUtils;
import org.nd4j.common.config.ND4JSystemProperties;
import org.nd4j.common.primitives.Counter;
import org.nd4j.linalg.api.buffer.DataBuffer;
import org.nd4j.linalg.api.memory.Deallocatable;
import org.nd4j.linalg.factory.Nd4j;

import java.lang.ref.ReferenceQueue;
import java.time.Duration;
import java.time.Instant;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedDeque;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Handles deallocation of offheap references.
 * This service starts a set of gc threads
 * and reference queues.
 * These reference queues are passed references to the jvm.
 *
 * The instances passed are of {@link Deallocatable}
 * These are all pointers with a deallocate() method which will be
 * called before being finally garbage collected by the JVM.
 *
 * The deallocator service handles generating unique ids by using an
 * atomic counter + an offset generated by {@link Deallocatable#getUniqueId()}
 *
 * Note that some buffers also may not be deallocated. Before deallocation a buffer
 * is checked for whether it's constant with {@link Deallocator#isConstant()}
 * If a reference is constant (this is usually only applicable to shape buffers)
 * then it will not be deallocated.
 *
 * You may turn the deallocator service off by setting {@link #noPointerGc}
 * to true with the system property: {@link ND4JSystemProperties#NO_ARRAY_GC}
 *
 * Note the service also checks "org.bytedeco.javacpp.noPointerGC"
 *
 * Also of note is you may access the references with {@link #getReferenceMap()}
 * DO THIS AT YOUR OWN RISK. ONLY FOR advanced users.
 *
 *
 * Known limitations:
 *
 * Due to relying on the JVM to supply garbage collectable references,
 * the allocator fall behind in deallocation. This happens VERY rarely and only at certain scales.
 *
 */
@Slf4j
public class DeallocatorService {
    private Thread[] deallocatorThreads;
    private ReferenceQueue<Deallocatable>[] queues;
    //note we do this for thread safety while reducing the amount of memory
    //strings take up. There is a performance trade off with respect to
    //ConcurrentHashMap (the prior implementation) which has higher throughput overall but is inefficient
    //for the amount of memory overhead it has. String compression
    //with a large number of objects is more important over throughput.
    @Getter
    private Map<Long,DeallocatableReference> referenceMap = Collections.synchronizedMap(new WeakHashMap<>());

    @Getter
    private Map<Long,String> referenceTypes = new ConcurrentHashMap<>();

    @Getter
    private Counter<String> allocated = new Counter<>();
    @Setter
    private Counter<String> deallocated = new Counter<>();

    private static AtomicBoolean blockDeallocator = new AtomicBoolean(false);

    private List<List<ReferenceQueue<Deallocatable>>> deviceMap = new ArrayList<>();
    private Boolean noPointerGc;
    private  int numThreads =  Integer.parseInt(System.getProperty(ND4JSystemProperties.DEALLOCATOR_SERVICE_GC_THREADS,"1"));

    private final transient AtomicLong counter = new AtomicLong(0);

    // Time-series tracking fields - only active when functrace is enabled
    private volatile boolean timeSeriesTrackingEnabled = false;
    private static final int MAX_SNAPSHOT_HISTORY = 1000;
    private final Deque<MemorySnapshot> snapshotHistory = new ConcurrentLinkedDeque<>();
    private final AtomicLong totalAllocations = new AtomicLong(0);
    private final AtomicLong totalDeallocations = new AtomicLong(0);
    private final AtomicLong totalBytesAllocated = new AtomicLong(0);
    private final AtomicLong totalBytesDeallocated = new AtomicLong(0);
    private final AtomicLong peakLiveCount = new AtomicLong(0);
    private final AtomicLong peakBytes = new AtomicLong(0);
    private final AtomicReference<Instant> trackingStartTime = new AtomicReference<>(null);

    /**
     * Represents a memory snapshot at a point in time for time-series tracking.
     * Used when functrace is enabled to track memory growth/shrinkage over time.
     */
    @Data
    @AllArgsConstructor
    public static class MemorySnapshot {
        private final Instant timestamp;
        private final long liveCount;
        private final long bytesInUse;
        private final long totalAllocations;
        private final long totalDeallocations;
        private final long totalBytesAllocated;
        private final long totalBytesDeallocated;
    }

    /**
     * Represents the delta between two memory snapshots.
     */
    @Data
    public static class MemoryDelta {
        private final Duration timeDelta;
        private final long liveCountDelta;
        private final long bytesDelta;
        private final long allocationsDelta;
        private final long deallocationsDelta;
        private final double allocationsPerSecond;
        private final double deallocationsPerSecond;
        private final double bytesPerSecond;

        public boolean isGrowing() {
            return bytesDelta > 0 || liveCountDelta > 0;
        }

        public boolean isShrinking() {
            return bytesDelta < 0 || liveCountDelta < 0;
        }

        public boolean isPotentialLeak() {
            // Consider it a potential leak if allocations significantly exceed deallocations
            return allocationsDelta > 0 && (deallocationsDelta == 0 ||
                   (double)allocationsDelta / Math.max(1, deallocationsDelta) > 1.5);
        }
    }


    /**
     * A listener for custom times for deallocation.
     * THe main use case is to prevent deallocation race conditions
     * when exiting a JVM process during tests.
     * Users may have other use cases for this listener, however.
     */
    public  interface CustomDeallocatorListener {

        void registerDataBuffer(DataBuffer reference);

        void registerDeallocatable(DeallocatableReference reference);
        /**
         * Adds a listener for deallocation.
         * This intercepts deallocate calls and calls them when the user is ready.
         * @param reference
         */
        void addForDeallocation(DeallocatableReference reference);
    }


    public void registerDataBufferToListener(DataBuffer reference) {
        for(CustomDeallocatorListener listener : listeners) {
            listener.registerDataBuffer(reference);
        }
    }
    public void registerDeallocatbleToListener(DeallocatableReference reference) {
        for(CustomDeallocatorListener listener : listeners) {
            listener.registerDeallocatable(reference);
        }
    }

    @Getter
    private List<CustomDeallocatorListener> listeners = new ArrayList<>();

    public DeallocatorService() {
        // we need to have at least 2 threads, but for CUDA we'd need at least numDevices threads, due to thread->device affinity
        int numDevices = Nd4j.getAffinityManager().getNumberOfDevices();

        for (int e = 0; e < numDevices; e++)
            deviceMap.add(new ArrayList<>());

        deallocatorThreads = new Thread[numThreads];
        queues = new ReferenceQueue[numThreads];
        noPointerGc = Boolean.parseBoolean(System.getProperty(ND4JSystemProperties.NO_ARRAY_GC,"false")) || Boolean.parseBoolean(System.getProperty("org.bytedeco.javacpp.nopointergc","false"));
        if(!noPointerGc) {
            for (int e = 0; e < numThreads; e++) {
                log.trace("Starting deallocator thread {}", e + 1);
                queues[e] = new ReferenceQueue<>();

                int deviceId = e % numDevices;
                // attaching queue to its own thread
                deallocatorThreads[e] = new DeallocatorServiceThread(queues[e], e, deviceId);
                deallocatorThreads[e].setName("DeallocatorServiceThread_" + e);
                deallocatorThreads[e].setDaemon(true);

                deviceMap.get(deviceId).add(queues[e]);

                deallocatorThreads[e].start();
            }
        } else {
            log.warn("Disabling automatic garbage collection since the system property " + ND4JSystemProperties.NO_ARRAY_GC + " or " + " org.bytedeco.javacpp.nopointergc was set to false");
        }

        // Initialize time-series tracking if functrace is enabled
        initializeTimeSeriesTracking();
    }

    /**
     * Initializes time-series tracking if functrace is enabled at build time.
     * This allows tracking memory growth/shrinkage over time, which is then
     * reported alongside other lifecycle tracker data.
     */
    private void initializeTimeSeriesTracking() {
        try {
            if (Nd4j.getNativeOps().isFuncTrace()) {
                timeSeriesTrackingEnabled = true;
                trackingStartTime.set(Instant.now());
                log.info("DeallocatorService time-series tracking enabled (functrace build detected)");
            }
        } catch (Exception e) {
            // NativeOps may not be initialized yet - that's OK, tracking stays disabled
            log.trace("Could not check isFuncTrace() status: {}", e.getMessage());
        }
    }

    /**
     * Returns whether time-series tracking is currently enabled.
     * @return true if tracking is enabled
     */
    public boolean isTimeSeriesTrackingEnabled() {
        return timeSeriesTrackingEnabled;
    }


    public void addListener(CustomDeallocatorListener listener) {
        listeners.add(listener);
    }

    /**
     * Returns the retained size
     * for each possible pointer type allocated.
     * @return
     */
    public String printRetainedSizesForTypes() {
        StringBuilder sb = new StringBuilder();
        for(String s : allocated.keySet()) {
            sb.append(s).append(" ").append(retainedSizeForType(s)).append("\n");
        }

        return sb.toString();
    }

    /**
     * Returns the allocated size for each pointer type
     * @return
     */
    public String allocatedTypes() {
        StringBuilder sb = new StringBuilder();
        for(String s : allocated.keySet()) {
            sb.append(s).append(" ").append(allocated).append("\n");
        }

        return sb.toString();
    }


    /**
     * Returns the retained size (the allocated size - deallocated size)
     * for each pointer type
     * @param pointer the type of pointer to get the retained size for
     * @return the retained size for the given pointer type
     */
    public long retainedSizeForType(String pointer) {
        return (long) (allocated.getCount(pointer) - deallocated.getCount(pointer));
    }


    public void toggleDeallocationBlock(boolean shouldBlock) {
        blockDeallocator.set(shouldBlock);
    }


    public long nextValue() {
        return counter.incrementAndGet();
    }

    /**
     * This method adds Deallocatable object instance to tracking system
     *
     * @param deallocatable object to track
     */
    public long pickObject(@NonNull Deallocatable deallocatable) {
        if(!noPointerGc) {

            val desiredDevice = deallocatable.targetDevice();
            val map = deviceMap.get(desiredDevice);


            val reference = new DeallocatableReference(deallocatable, map.get(RandomUtils.nextInt(0, numThreads)));
            referenceMap.put(deallocatable.getUniqueId(), reference);
            return deallocatable.getUniqueId();
        }

        return -1;
    }


    private class DeallocatorServiceThread extends Thread implements Runnable {
        private final ReferenceQueue<Deallocatable> queue;
        private final int threadIdx;
        public static final String DeallocatorThreadNamePrefix = "DeallocatorServiceThread thread ";
        private final int deviceId;

        private DeallocatorServiceThread(@NonNull ReferenceQueue<Deallocatable> queue, int threadIdx, int deviceId) {
            this.queue = queue;
            this.threadIdx = threadIdx;
            this.setName(DeallocatorThreadNamePrefix + threadIdx);
            this.deviceId = deviceId;
            setContextClassLoader(null);
        }

        @SneakyThrows
        @Override
        public void run() {
            Nd4j.getAffinityManager().unsafeSetDevice(deviceId);
            boolean canRun = true;
            while (canRun) {
                while(blockDeallocator.get()) {
                    Thread.sleep(1000);
                }
                // if periodicGc is enabled, only first thread will call for it
                if (threadIdx == 0 && Nd4j.getMemoryManager().getAutoGcWindow() > 0) {
                    val reference = (DeallocatableReference) queue.poll();
                    if (reference == null) {
                        val timeout = Nd4j.getMemoryManager().getAutoGcWindow();
                        try {
                            Thread.sleep(timeout);
                            Nd4j.getMemoryManager().invokeGc();
                        } catch (InterruptedException e) {
                            canRun = false;
                        }
                    } else {
                        // invoking deallocator
                        if (reference != null) {
                            if(!listeners.isEmpty()) {
                                reference.deallocate();
                                if(referenceMap.containsKey(reference.getId()))
                                    referenceMap.remove(reference.getId());
                            }

                            else {
                                for(CustomDeallocatorListener listener : listeners)
                                    listener.addForDeallocation(reference);
                            }


                        }
                    }
                } else {
                    try {
                        val reference = (DeallocatableReference) queue.remove();
                        if (reference == null)
                            continue;

                        if(!listeners.isEmpty()) {
                            reference.deallocate();
                            if(referenceMap.containsKey(reference.getId()))
                                referenceMap.remove(reference.getId());

                        }

                        else {
                            for(CustomDeallocatorListener listener : listeners)
                                listener.addForDeallocation(reference);
                        }

                    } catch (InterruptedException e) {
                        canRun = false;
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                }
            }
        }
    }

    // ===============================
    // Time-Series Tracking Methods
    // ===============================

    /**
     * Records an allocation event for time-series tracking.
     * Only tracks if functrace is enabled.
     *
     * @param bytes the number of bytes allocated
     */
    public void recordAllocation(long bytes) {
        if (!timeSeriesTrackingEnabled) return;

        totalAllocations.incrementAndGet();
        totalBytesAllocated.addAndGet(bytes);

        // Update peak tracking
        long currentLive = totalAllocations.get() - totalDeallocations.get();
        peakLiveCount.updateAndGet(peak -> Math.max(peak, currentLive));

        long currentBytes = totalBytesAllocated.get() - totalBytesDeallocated.get();
        peakBytes.updateAndGet(peak -> Math.max(peak, currentBytes));
    }

    /**
     * Records a deallocation event for time-series tracking.
     * Only tracks if functrace is enabled.
     *
     * @param bytes the number of bytes deallocated
     */
    public void recordDeallocation(long bytes) {
        if (!timeSeriesTrackingEnabled) return;

        totalDeallocations.incrementAndGet();
        totalBytesDeallocated.addAndGet(bytes);
    }

    /**
     * Takes a snapshot of current memory state.
     * Only active if functrace is enabled.
     *
     * @return the current memory snapshot, or null if tracking is disabled
     */
    public MemorySnapshot takeSnapshot() {
        if (!timeSeriesTrackingEnabled) return null;

        long live = totalAllocations.get() - totalDeallocations.get();
        long bytesInUse = totalBytesAllocated.get() - totalBytesDeallocated.get();

        MemorySnapshot snapshot = new MemorySnapshot(
                Instant.now(),
                live,
                bytesInUse,
                totalAllocations.get(),
                totalDeallocations.get(),
                totalBytesAllocated.get(),
                totalBytesDeallocated.get()
        );

        // Store in history, maintaining max size
        snapshotHistory.addLast(snapshot);
        while (snapshotHistory.size() > MAX_SNAPSHOT_HISTORY) {
            snapshotHistory.removeFirst();
        }

        return snapshot;
    }

    /**
     * Calculates the delta between two snapshots.
     *
     * @param older the older snapshot
     * @param newer the newer snapshot
     * @return the delta between the snapshots
     */
    public MemoryDelta calculateDelta(MemorySnapshot older, MemorySnapshot newer) {
        Duration timeDelta = Duration.between(older.getTimestamp(), newer.getTimestamp());
        long seconds = Math.max(1, timeDelta.getSeconds());

        long liveCountDelta = newer.getLiveCount() - older.getLiveCount();
        long bytesDelta = newer.getBytesInUse() - older.getBytesInUse();
        long allocationsDelta = newer.getTotalAllocations() - older.getTotalAllocations();
        long deallocationsDelta = newer.getTotalDeallocations() - older.getTotalDeallocations();

        return new MemoryDelta(
                timeDelta,
                liveCountDelta,
                bytesDelta,
                allocationsDelta,
                deallocationsDelta,
                (double) allocationsDelta / seconds,
                (double) deallocationsDelta / seconds,
                (double) bytesDelta / seconds
        );
    }

    /**
     * Gets the snapshot history for analysis.
     *
     * @return list of snapshots, oldest first
     */
    public List<MemorySnapshot> getSnapshotHistory() {
        return new ArrayList<>(snapshotHistory);
    }

    /**
     * Detects potential memory leak by analyzing the growth pattern.
     *
     * @return true if a potential leak is detected
     */
    public boolean detectPotentialLeak() {
        if (!timeSeriesTrackingEnabled || snapshotHistory.size() < 2) return false;

        MemorySnapshot oldest = snapshotHistory.peekFirst();
        MemorySnapshot newest = snapshotHistory.peekLast();

        if (oldest == null || newest == null) return false;

        MemoryDelta delta = calculateDelta(oldest, newest);
        return delta.isPotentialLeak();
    }

    /**
     * Gets a summary of the current tracking state.
     * This data is meant to be pushed to the C++ lifecycle tracker.
     *
     * @return summary map suitable for native integration
     */
    public Map<String, Long> getTrackingSummary() {
        Map<String, Long> summary = new LinkedHashMap<>();
        if (!timeSeriesTrackingEnabled) {
            summary.put("enabled", 0L);
            return summary;
        }

        summary.put("enabled", 1L);
        summary.put("totalAllocations", totalAllocations.get());
        summary.put("totalDeallocations", totalDeallocations.get());
        summary.put("totalBytesAllocated", totalBytesAllocated.get());
        summary.put("totalBytesDeallocated", totalBytesDeallocated.get());
        summary.put("currentLiveCount", totalAllocations.get() - totalDeallocations.get());
        summary.put("currentBytesInUse", totalBytesAllocated.get() - totalBytesDeallocated.get());
        summary.put("peakLiveCount", peakLiveCount.get());
        summary.put("peakBytes", peakBytes.get());
        summary.put("snapshotCount", (long) snapshotHistory.size());

        Instant start = trackingStartTime.get();
        if (start != null) {
            summary.put("trackingDurationMs", Duration.between(start, Instant.now()).toMillis());
        }

        return summary;
    }

    /**
     * Gets the peak live count observed since tracking started.
     */
    public long getPeakLiveCount() {
        return peakLiveCount.get();
    }

    /**
     * Gets the peak bytes in use observed since tracking started.
     */
    public long getPeakBytes() {
        return peakBytes.get();
    }

    /**
     * Gets the tracking duration since initialization.
     */
    public Duration getTrackingDuration() {
        Instant start = trackingStartTime.get();
        if (start == null) return Duration.ZERO;
        return Duration.between(start, Instant.now());
    }

    /**
     * Pushes current tracking state to native lifecycle tracker.
     * This should be called periodically to sync Java-side tracking
     * with the C++ UnifiedMemoryReporter.
     */
    public void pushToNativeTracker() {
        if (!timeSeriesTrackingEnabled) return;

        try {
            Map<String, Long> summary = getTrackingSummary();
            Nd4j.getNativeOps().recordDeallocatorServiceSnapshot(
                    summary.get("totalAllocations"),
                    summary.get("totalDeallocations"),
                    summary.get("totalBytesAllocated"),
                    summary.get("totalBytesDeallocated"),
                    summary.get("peakLiveCount"),
                    summary.get("peakBytes")
            );
        } catch (Exception e) {
            log.trace("Failed to push to native tracker: {}", e.getMessage());
        }
    }

    /**
     * Resets all time-series tracking counters.
     * Useful for testing or after a memory checkpoint.
     */
    public void resetTracking() {
        if (!timeSeriesTrackingEnabled) return;

        totalAllocations.set(0);
        totalDeallocations.set(0);
        totalBytesAllocated.set(0);
        totalBytesDeallocated.set(0);
        peakLiveCount.set(0);
        peakBytes.set(0);
        snapshotHistory.clear();
        trackingStartTime.set(Instant.now());
    }

}
