package bench;

import java.io.BufferedWriter;
import java.io.IOException;
import java.lang.management.ManagementFactory;
import java.lang.management.MemoryMXBean;
import java.lang.management.MemoryPoolMXBean;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.util.List;

/**
 * Memory sampling for the MegaMek density benchmark.
 *
 * <p>Three numbers matter and they are not interchangeable:
 * <ul>
 *   <li>{@code rssKb} - what the OS actually charges us, and therefore what sizes a VM.</li>
 *   <li>{@code heapUsedKb} - live + garbage. Noisy; only meaningful right after a GC.</li>
 *   <li>{@code liveHeapKb} - taken after forced collection. This is the real retained
 *       footprint, and the only honest basis for a "marginal cost per game" claim.</li>
 * </ul>
 */
public final class Sampler {

    private final Path csv;
    private final long startNanos = System.nanoTime();
    private final MemoryMXBean memoryBean = ManagementFactory.getMemoryMXBean();

    public Sampler(Path csv) throws IOException {
        this.csv = csv;
        Files.createDirectories(csv.toAbsolutePath().getParent());
        Files.writeString(csv,
              "elapsedSec,label,gamesStarted,gamesLive,totalRounds,rssKb,heapUsedKb,"
                    + "heapCommittedKb,nonHeapUsedKb,metaspaceKb,liveHeapKb\n",
              StandardCharsets.UTF_8);
    }

    /** Resident set size from /proc/self/status; -1 where unavailable. */
    public static long rssKb() {
        try {
            for (String line : Files.readAllLines(Path.of("/proc/self/status"))) {
                if (line.startsWith("VmRSS:")) {
                    return Long.parseLong(line.replaceAll("[^0-9]", ""));
                }
            }
        } catch (Exception ignored) {
            // fall through
        }
        return -1L;
    }

    private static long metaspaceKb() {
        long total = 0L;
        List<MemoryPoolMXBean> pools = ManagementFactory.getMemoryPoolMXBeans();
        for (MemoryPoolMXBean pool : pools) {
            String name = pool.getName();
            if (name.contains("Metaspace") || name.contains("Class space")) {
                total += pool.getUsage().getUsed();
            }
        }
        return total / 1024L;
    }

    /**
     * Forces collection and returns retained heap.
     *
     * <p>One System.gc() is not enough under G1: weakly reachable objects and
     * humongous regions routinely survive the first pass. Three passes with a
     * pause between them gets close enough to the true live set that the
     * per-game delta is trustworthy.
     */
    public static long liveHeapKb() {
        for (int i = 0; i < 3; i++) {
            System.gc();
            try {
                Thread.sleep(250L);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                break;
            }
        }
        Runtime rt = Runtime.getRuntime();
        return (rt.totalMemory() - rt.freeMemory()) / 1024L;
    }

    /** Cheap sample: no forced GC, safe to call on a tight polling loop. */
    public void sample(String label, int gamesStarted, int gamesLive, int totalRounds) {
        write(label, gamesStarted, gamesLive, totalRounds, -1L);
    }

    /** Expensive sample: forces GC first so liveHeapKb is meaningful. */
    public void checkpoint(String label, int gamesStarted, int gamesLive, int totalRounds) {
        write(label, gamesStarted, gamesLive, totalRounds, liveHeapKb());
    }

    private void write(String label, int gamesStarted, int gamesLive, int totalRounds, long liveHeapKb) {
        double elapsedSec = (System.nanoTime() - startNanos) / 1_000_000_000.0;
        long heapUsedKb = memoryBean.getHeapMemoryUsage().getUsed() / 1024L;
        long heapCommittedKb = memoryBean.getHeapMemoryUsage().getCommitted() / 1024L;
        long nonHeapUsedKb = memoryBean.getNonHeapMemoryUsage().getUsed() / 1024L;

        String row = String.format("%.1f,%s,%d,%d,%d,%d,%d,%d,%d,%d,%d%n",
              elapsedSec, label, gamesStarted, gamesLive, totalRounds,
              rssKb(), heapUsedKb, heapCommittedKb, nonHeapUsedKb, metaspaceKb(), liveHeapKb);

        try (BufferedWriter out = Files.newBufferedWriter(csv, StandardCharsets.UTF_8,
              StandardOpenOption.APPEND)) {
            out.write(row);
        } catch (IOException e) {
            System.err.println("[sampler] write failed: " + e.getMessage());
        }

        if (liveHeapKb >= 0) {
            System.out.printf("[%6.1fs] %-22s games=%d live=%d rounds=%d  rss=%dMB  liveHeap=%dMB  meta=%dMB%n",
                  elapsedSec, label, gamesStarted, gamesLive, totalRounds,
                  rssKb() / 1024, liveHeapKb / 1024, metaspaceKb() / 1024);
        }
    }
}