package bench;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.apache.logging.log4j.core.LoggerContext;
import org.apache.logging.log4j.core.config.Configurator;

import java.net.URI;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

public class BenchRunner {
    private final BenchConfig config;

    public BenchRunner(BenchConfig config) {
        this.config = config;
    }

    public BenchResult run(String configName) throws Exception {
        dropCaches();
        reconfigure(configName);

        Logger logger = LogManager.getLogger("bench");
        String message = config.generateMessage();
        int messageBytes = message.getBytes().length;

        CountDownLatch startLatch = new CountDownLatch(1);
        AtomicBoolean running = new AtomicBoolean(true);
        AtomicLong eventCounter = new AtomicLong(0);
        AtomicLong targetEvents = new AtomicLong(config.getTotalEvents());

        List<Thread> threads = new ArrayList<>();
        for (int i = 0; i < config.getThreads(); i++) {
            Thread t = new Thread(new LogWorker(
                logger, message, startLatch, running, eventCounter, targetEvents, config.getMode()
            ));
            t.start();
            threads.add(t);
        }

        // Warmup phase
        if (config.getWarmupSeconds() > 0) {
            startLatch.countDown();
            Thread.sleep(config.getWarmupSeconds() * 1000);
            eventCounter.set(0);
        } else {
            startLatch.countDown();
        }

        long startTime = System.nanoTime();

        if (config.getMode() == BenchConfig.Mode.DURATION) {
            Thread.sleep(config.getDurationSeconds() * 1000);
            running.set(false);
        } else {
            while (eventCounter.get() < targetEvents.get()) {
                Thread.sleep(10);
            }
            running.set(false);
        }

        long endTime = System.nanoTime();

        for (Thread t : threads) {
            t.join(5000);
        }

        long events = eventCounter.get();
        double durationSec = (endTime - startTime) / 1_000_000_000.0;
        double eventsPerSec = events / durationSec;
        double bytesPerSec = (events * messageBytes) / durationSec;
        double mbPerSec = bytesPerSec / (1024 * 1024);

        // Shutdown log4j context to flush and release resources
        LoggerContext ctx = (LoggerContext) LogManager.getContext(false);
        ctx.stop();

        return new BenchResult(configName, config.getThreads(), events, durationSec, eventsPerSec, mbPerSec);
    }

    private void dropCaches() {
        try {
            ProcessBuilder pb = new ProcessBuilder("sh", "-c", "sync; echo 3 > /proc/sys/vm/drop_caches");
            pb.inheritIO();
            Process p = pb.start();
            int exitCode = p.waitFor();
            if (exitCode != 0) {
                System.err.println("Warning: Failed to drop caches (exit code " + exitCode + "). Run as root for accurate benchmarks.");
            }
        } catch (Exception e) {
            System.err.println("Warning: Could not drop caches: " + e.getMessage());
        }
    }

    private void reconfigure(String configName) throws Exception {
        // Fully shutdown existing context
        LoggerContext ctx = (LoggerContext) LogManager.getContext(false);
        ctx.stop();
        
        // Set async ring buffer size via system property (must be set before reconfigure)
        if (configName.startsWith("async-")) {
            String size = configName.substring(6); // e.g., "1k", "4k", "10k"
            int bufferSize = switch (size) {
                case "1k" -> 1024;
                case "4k" -> 4096;
                case "10k" -> 10240;
                default -> 4096;
            };
            System.setProperty("log4j2.asyncLoggerRingBufferSize", String.valueOf(bufferSize));
            System.setProperty("log4j2.contextSelector", "org.apache.logging.log4j.core.async.AsyncLoggerContextSelector");
        } else {
            System.clearProperty("log4j2.asyncLoggerRingBufferSize");
            System.clearProperty("log4j2.contextSelector");
        }

        String resourcePath = "log4j2-" + configName + ".xml";
        URI uri = getClass().getClassLoader().getResource(resourcePath).toURI();
        Configurator.reconfigure(uri);
    }

    public static class BenchResult {
        public final String configName;
        public final int threads;
        public final long events;
        public final double durationSec;
        public final double eventsPerSec;
        public final double mbPerSec;

        public BenchResult(String configName, int threads, long events,
                          double durationSec, double eventsPerSec, double mbPerSec) {
            this.configName = configName;
            this.threads = threads;
            this.events = events;
            this.durationSec = durationSec;
            this.eventsPerSec = eventsPerSec;
            this.mbPerSec = mbPerSec;
        }

        public String toCsv() {
            return String.format("%s,%d,%d,%.2f,%.0f,%.2f",
                configName, threads, events, durationSec, eventsPerSec, mbPerSec);
        }

        @Override
        public String toString() {
            return String.format("%-16s | %3d threads | %,12d events | %.2fs | %,.0f events/s | %.2f MB/s",
                configName, threads, events, durationSec, eventsPerSec, mbPerSec);
        }
    }
}