package ai.nektron.grownet.debug;

import java.io.PrintWriter;
import java.io.Writer;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;

public final class BenchmarkRunner {
    private static final List<String> FIELD_ORDER = List.of(
            "lang",
            "scenario",
            "ticks",
            "warmup",
            "iterations",
            "repeat",
            "trace",
            "run",
            "count",
            "runs",
            "min_us",
            "p50_us",
            "p95_us",
            "p99_us",
            "max_us",
            "mean_us",
            "overall_mean_us",
            "total_elapsed_ms",
            "overall_elapsed_ms",
            "throughput_per_sec",
            "best_p50_us",
            "best_p99_us",
            "best_throughput_per_sec"
    );

    private BenchmarkRunner() {
    }

    public static void run(TraceConfig traceConfig, Writer benchmarkWriter, Writer measuredTraceWriter) throws Exception {
        PrintWriter printWriter = new PrintWriter(benchmarkWriter, true);
        emitStart(traceConfig, printWriter);

        List<BenchmarkStats> runStats = new ArrayList<>();
        for (int runOrdinal = 1; runOrdinal <= traceConfig.getBenchmarkRepeat(); runOrdinal += 1) {
            warmup(traceConfig);
            BenchmarkStats stats = measure(traceConfig, measuredTraceWriter);
            runStats.add(stats);
            emitResult(traceConfig, runOrdinal, stats, printWriter);
        }

        emitSummary(traceConfig, summarize(runStats), printWriter);
        printWriter.flush();
    }

    private static void warmup(TraceConfig traceConfig) throws Exception {
        Writer discardWriter = new DiscardingWriter();
        for (int warmupOrdinal = 0; warmupOrdinal < traceConfig.getBenchmarkWarmup(); warmupOrdinal += 1) {
            GrowNetTraceMain.executeOnce(traceConfig, discardWriter);
        }
    }

    private static BenchmarkStats measure(TraceConfig traceConfig, Writer measuredTraceWriter) throws Exception {
        long[] iterationDurationsNanos = new long[traceConfig.getBenchmarkIterations()];
        long benchmarkStartNanos = System.nanoTime();
        for (int iterationOrdinal = 0; iterationOrdinal < traceConfig.getBenchmarkIterations(); iterationOrdinal += 1) {
            long iterationStartNanos = System.nanoTime();
            GrowNetTraceMain.executeOnce(traceConfig, measuredTraceWriter);
            iterationDurationsNanos[iterationOrdinal] = System.nanoTime() - iterationStartNanos;
        }
        long benchmarkElapsedNanos = System.nanoTime() - benchmarkStartNanos;
        return BenchmarkStats.fromIterationNanos(iterationDurationsNanos, benchmarkElapsedNanos);
    }

    private static void emitStart(TraceConfig traceConfig, PrintWriter printWriter) {
        Map<String, Object> eventFields = new LinkedHashMap<>();
        eventFields.put("lang", traceConfig.getLang());
        eventFields.put("scenario", traceConfig.getScenario());
        eventFields.put("ticks", traceConfig.getTicks());
        eventFields.put("warmup", traceConfig.getBenchmarkWarmup());
        eventFields.put("iterations", traceConfig.getBenchmarkIterations());
        eventFields.put("repeat", traceConfig.getBenchmarkRepeat());
        eventFields.put("trace", traceConfig.isTraceEnabled());
        printWriter.println(formatLine(traceConfig, "benchmark.start", eventFields));
    }

    private static void emitResult(
            TraceConfig traceConfig,
            int runOrdinal,
            BenchmarkStats stats,
            PrintWriter printWriter) {
        Map<String, Object> eventFields = new LinkedHashMap<>();
        eventFields.put("lang", traceConfig.getLang());
        eventFields.put("scenario", traceConfig.getScenario());
        eventFields.put("run", runOrdinal);
        eventFields.put("count", stats.getCount());
        eventFields.put("min_us", stats.getMinUs());
        eventFields.put("p50_us", stats.getP50Us());
        eventFields.put("p95_us", stats.getP95Us());
        eventFields.put("p99_us", stats.getP99Us());
        eventFields.put("max_us", stats.getMaxUs());
        eventFields.put("mean_us", stats.getMeanUs());
        eventFields.put("total_elapsed_ms", stats.getTotalElapsedMs());
        eventFields.put("throughput_per_sec", stats.getThroughputPerSec());
        printWriter.println(formatLine(traceConfig, "benchmark.result", eventFields));
    }

    private static void emitSummary(
            TraceConfig traceConfig,
            BenchmarkSummary summary,
            PrintWriter printWriter) {
        Map<String, Object> eventFields = new LinkedHashMap<>();
        eventFields.put("lang", traceConfig.getLang());
        eventFields.put("scenario", traceConfig.getScenario());
        eventFields.put("runs", summary.runs);
        eventFields.put("overall_mean_us", summary.overallMeanUs);
        eventFields.put("overall_elapsed_ms", summary.overallElapsedMs);
        eventFields.put("best_p50_us", summary.bestP50Us);
        eventFields.put("best_p99_us", summary.bestP99Us);
        eventFields.put("best_throughput_per_sec", summary.bestThroughputPerSec);
        printWriter.println(formatLine(traceConfig, "benchmark.summary", eventFields));
    }

    private static String formatLine(TraceConfig traceConfig, String eventName, Map<String, Object> eventFields) {
        return "jsonl".equals(traceConfig.getBenchmarkFormat())
                ? formatJsonLine(eventName, eventFields)
                : formatTextLine(eventName, eventFields);
    }

    private static String formatTextLine(String eventName, Map<String, Object> eventFields) {
        StringBuilder builder = new StringBuilder();
        builder.append('[').append(eventName).append(']');
        for (String fieldName : orderedKeys(eventFields)) {
            Object fieldValue = eventFields.get(fieldName);
            if (fieldValue == null) {
                continue;
            }
            builder.append(' ')
                    .append(fieldName)
                    .append('=')
                    .append(formatScalar(fieldValue));
        }
        return builder.toString();
    }

    private static String formatJsonLine(String eventName, Map<String, Object> eventFields) {
        StringBuilder builder = new StringBuilder();
        builder.append('{');
        builder.append("\"event\":\"").append(escapeJson(eventName)).append('"');
        for (String fieldName : orderedKeys(eventFields)) {
            Object fieldValue = eventFields.get(fieldName);
            if (fieldValue == null) {
                continue;
            }
            builder.append(',');
            builder.append('"').append(escapeJson(fieldName)).append('"').append(':');
            appendJsonValue(builder, fieldValue);
        }
        builder.append('}');
        return builder.toString();
    }

    private static List<String> orderedKeys(Map<String, Object> eventFields) {
        List<String> orderedKeys = new ArrayList<>();
        for (String preferredField : FIELD_ORDER) {
            if (eventFields.containsKey(preferredField)) {
                orderedKeys.add(preferredField);
            }
        }
        List<String> additionalKeys = new ArrayList<>();
        for (String fieldName : eventFields.keySet()) {
            if (!FIELD_ORDER.contains(fieldName)) {
                additionalKeys.add(fieldName);
            }
        }
        additionalKeys.sort(Comparator.naturalOrder());
        orderedKeys.addAll(additionalKeys);
        return orderedKeys;
    }

    private static String formatScalar(Object fieldValue) {
        if (fieldValue instanceof Double doubleValue) {
            return String.format(Locale.ROOT, "%.3f", doubleValue);
        }
        if (fieldValue instanceof Float floatValue) {
            return String.format(Locale.ROOT, "%.3f", floatValue);
        }
        return String.valueOf(fieldValue);
    }

    private static void appendJsonValue(StringBuilder builder, Object fieldValue) {
        if (fieldValue instanceof Double doubleValue) {
            builder.append(String.format(Locale.ROOT, "%.3f", doubleValue));
            return;
        }
        if (fieldValue instanceof Float floatValue) {
            builder.append(String.format(Locale.ROOT, "%.3f", floatValue));
            return;
        }
        if (fieldValue instanceof Number || fieldValue instanceof Boolean) {
            builder.append(fieldValue);
            return;
        }
        builder.append('"').append(escapeJson(String.valueOf(fieldValue))).append('"');
    }

    private static String escapeJson(String rawValue) {
        StringBuilder builder = new StringBuilder();
        for (int charIndex = 0; charIndex < rawValue.length(); charIndex += 1) {
            char currentChar = rawValue.charAt(charIndex);
            switch (currentChar) {
                case '\\' -> builder.append("\\\\");
                case '"' -> builder.append("\\\"");
                case '\n' -> builder.append("\\n");
                case '\r' -> builder.append("\\r");
                case '\t' -> builder.append("\\t");
                default -> builder.append(currentChar);
            }
        }
        return builder.toString();
    }

    private static BenchmarkSummary summarize(List<BenchmarkStats> runStats) {
        long totalLatencyNanos = 0L;
        long totalElapsedNanos = 0L;
        long totalCount = 0L;
        double bestP50Us = Double.POSITIVE_INFINITY;
        double bestP99Us = Double.POSITIVE_INFINITY;
        double bestThroughputPerSec = 0.0;

        for (BenchmarkStats stats : runStats) {
            totalLatencyNanos += stats.getTotalLatencyNanos();
            totalElapsedNanos += stats.getTotalElapsedNanos();
            totalCount += stats.getCount();
            bestP50Us = Math.min(bestP50Us, stats.getP50Us());
            bestP99Us = Math.min(bestP99Us, stats.getP99Us());
            bestThroughputPerSec = Math.max(bestThroughputPerSec, stats.getThroughputPerSec());
        }

        double overallMeanUs = totalCount == 0L ? 0.0 : (totalLatencyNanos / 1_000.0) / totalCount;
        double overallElapsedMs = totalElapsedNanos / 1_000_000.0;
        return new BenchmarkSummary(
                runStats.size(),
                overallMeanUs,
                overallElapsedMs,
                bestP50Us == Double.POSITIVE_INFINITY ? 0.0 : bestP50Us,
                bestP99Us == Double.POSITIVE_INFINITY ? 0.0 : bestP99Us,
                bestThroughputPerSec
        );
    }

    private record BenchmarkSummary(
            int runs,
            double overallMeanUs,
            double overallElapsedMs,
            double bestP50Us,
            double bestP99Us,
            double bestThroughputPerSec) {
    }

    private static final class DiscardingWriter extends Writer {
        @Override
        public void write(char[] buffer, int offset, int length) {
        }

        @Override
        public void flush() {
        }

        @Override
        public void close() {
        }
    }
}
