package bench;

import picocli.CommandLine;
import picocli.CommandLine.Command;
import picocli.CommandLine.Option;

import java.io.FileWriter;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.Callable;

@Command(name = "log4jbench", mixinStandardHelpOptions = true, version = "1.0",
         description = "Benchmark Log4j2 logging throughput with various configurations")
public class Main implements Callable<Integer> {

    @Option(names = {"-t", "--threads"}, description = "Number of concurrent threads (default: available CPUs)")
    private int threads = Runtime.getRuntime().availableProcessors();

    @Option(names = {"-m", "--mode"}, description = "Test mode: duration or events (default: duration)")
    private String mode = "duration";

    @Option(names = {"-d", "--duration"}, description = "Test duration in seconds (default: 10)")
    private long duration = 10;

    @Option(names = {"-e", "--events"}, description = "Total events to log in events mode (default: 1000000)")
    private long events = 1_000_000;

    @Option(names = {"-w", "--warmup"}, description = "Warmup duration in seconds (default: 3)")
    private long warmup = 3;

    @Option(names = {"-s", "--msg-size"}, description = "Message size in characters (default: 100)")
    private int msgSize = 100;

    @Option(names = {"-c", "--configs"}, description = "Comma-separated list of configs to test (default: all)")
    private String configs = null;

    @Option(names = {"-o", "--output"}, description = "Output CSV file (default: stdout only)")
    private String outputFile = null;

    @Option(names = {"--flush"}, description = "Call LogManager.shutdown() to flush at end of each test (default: false)")
    private boolean flush = false;

    @Option(names = {"--drop-caches"}, description = "Drop OS file caches before each test (requires sudo, default: false)")
    private boolean dropCaches = false;

    public static void main(String[] args) {
        // Must set context selector before ANY Log4j initialization
        // We'll set it dynamically per-config in BenchRunner instead
        int exitCode = new CommandLine(new Main()).execute(args);
        System.exit(exitCode);
    }

    @Override
    public Integer call() throws Exception {
        BenchConfig config = new BenchConfig();
        config.setThreads(threads);
        config.setMode("events".equalsIgnoreCase(mode) ? BenchConfig.Mode.EVENTS : BenchConfig.Mode.DURATION);
        config.setDurationSeconds(duration);
        config.setTotalEvents(events);
        config.setWarmupSeconds(warmup);
        config.setMessageSize(msgSize);
        config.setOutputFile(outputFile);

        if (configs != null && !configs.isBlank()) {
            config.setConfigs(Arrays.asList(configs.split(",")));
        }

        System.out.println("=== Log4j2 Benchmark ===");
        System.out.printf("Threads: %d | Mode: %s | Message size: %d chars | flush: %s%n",
            config.getThreads(), config.getMode(), config.getMessageSize(), flush);
        if (config.getMode() == BenchConfig.Mode.DURATION) {
            System.out.printf("Duration: %ds | Warmup: %ds%n", config.getDurationSeconds(), config.getWarmupSeconds());
        } else {
            System.out.printf("Target events: %,d | Warmup: %ds%n", config.getTotalEvents(), config.getWarmupSeconds());
        }
        System.out.printf("Configs: %s%n", config.getConfigs());
        System.out.println();

        PrintWriter csvWriter = null;
        if (outputFile != null) {
            csvWriter = new PrintWriter(new FileWriter(outputFile));
            csvWriter.println("config,threads,events,duration_sec,events_per_sec,mb_per_sec");
        }

        String jarPath = Main.class.getProtectionDomain().getCodeSource().getLocation().toURI().getPath();
        String javaHome = System.getProperty("java.home");
        String javaBin = javaHome + "/bin/java";

        for (String configName : config.getConfigs()) {
            System.out.printf("Running: %s ...%n", configName);
            try {
                // Drop caches before each test if requested
                if (dropCaches) {
                    dropCaches();
                }

                // Build command with appropriate system properties
                List<String> cmd = new ArrayList<>();
                cmd.add(javaBin);
                
                // Set async properties if needed
                if (configName.startsWith("async-")) {
                    String size = configName.substring(6);
                    int bufferSize = switch (size) {
                        case "1k" -> 1024;
                        case "4k" -> 4096;
                        case "10k" -> 10240;
                        case "1m" -> 1048576;
                        default -> 4096;
                    };
                    cmd.add("-Dlog4j2.contextSelector=org.apache.logging.log4j.core.async.AsyncLoggerContextSelector");
                    cmd.add("-Dlog4j2.asyncLoggerRingBufferSize=" + bufferSize);
                }
                
                // Set log4j config file (use classpath reference)
                cmd.add("-Dlog4j2.configurationFile=classpath:log4j2-" + configName + ".xml");
                
                // Pass flush setting
                cmd.add("-Dbench.flush=" + flush);
                
                cmd.add("-cp");
                cmd.add(jarPath);
                cmd.add("bench.SingleBench");
                cmd.add(configName);
                cmd.add(String.valueOf(config.getThreads()));
                cmd.add(String.valueOf(config.getDurationSeconds()));
                cmd.add(String.valueOf(config.getWarmupSeconds()));
                cmd.add(String.valueOf(config.getMessageSize()));

                ProcessBuilder pb = new ProcessBuilder(cmd);
                pb.redirectErrorStream(true);
                Process p = pb.start();
                
                long pid = p.pid();
                System.out.printf("  [PID: %d]%n", pid);
                
                String output = new String(p.getInputStream().readAllBytes());
                int exitCode = p.waitFor();
                
                if (exitCode == 0 && !output.isBlank()) {
                    String[] parts = output.trim().split(",");
                    if (parts.length >= 6) {
                        String resultLine = String.format("%-16s | %3d threads | %,12d events | %.2fs | %,.0f events/s | %.2f MB/s",
                            parts[0], Integer.parseInt(parts[1]), Long.parseLong(parts[2]),
                            Double.parseDouble(parts[3]), Double.parseDouble(parts[4]), Double.parseDouble(parts[5]));
                        System.out.println(resultLine);
                        if (csvWriter != null) {
                            csvWriter.println(output.trim());
                            csvWriter.flush();
                        }
                    } else {
                        System.out.println(output);
                    }
                } else {
                    System.err.printf("Error (exit %d): %s%n", exitCode, output);
                }
            } catch (Exception e) {
                System.err.printf("Error running config %s: %s%n", configName, e.getMessage());
                e.printStackTrace();
            }
            System.out.println();
        }

        if (csvWriter != null) {
            csvWriter.close();
            System.out.printf("Results written to: %s%n", outputFile);
        }

        return 0;
    }

    private void dropCaches() {
        try {
            System.out.print("  Dropping caches... ");
            System.out.flush();
            ProcessBuilder pb = new ProcessBuilder("sh", "-c", "sync; echo 3 | sudo tee /proc/sys/vm/drop_caches > /dev/null");
            Process p = pb.start();
            int exitCode = p.waitFor();
            if (exitCode == 0) {
                System.out.println("done");
            } else {
                System.out.println("failed (exit code " + exitCode + "). Ensure sudo is available.");
            }
        } catch (Exception e) {
            System.out.println("failed: " + e.getMessage());
        }
    }
}