package ai.nektron.grownet.debug;

import java.nio.file.Path;
import java.util.Collections;
import java.util.LinkedHashSet;
import java.util.Locale;
import java.util.Set;

public final class TraceConfig {
    private final boolean traceEnabled;
    private final boolean debugEnabled;
    private final boolean benchmarkEnabled;
    private final String lang;
    private final String traceFormat;
    private final String benchmarkFormat;
    private final Set<String> traceFamilies;
    private final boolean traceSignal;
    private final String scenario;
    private final double value;
    private final int ticks;
    private final String outTarget;
    private final int benchmarkWarmup;
    private final int benchmarkIterations;
    private final int benchmarkRepeat;
    private final String benchmarkOutputTarget;
    private final String traceNeuron;
    private final Integer traceLayer;
    private final boolean helpRequested;

    private TraceConfig(
            boolean traceEnabled,
            boolean debugEnabled,
            boolean benchmarkEnabled,
            String lang,
            String traceFormat,
            String benchmarkFormat,
            Set<String> traceFamilies,
            boolean traceSignal,
            String scenario,
            double value,
            int ticks,
            String outTarget,
            int benchmarkWarmup,
            int benchmarkIterations,
            int benchmarkRepeat,
            String benchmarkOutputTarget,
            String traceNeuron,
            Integer traceLayer,
            boolean helpRequested) {
        this.traceEnabled = traceEnabled;
        this.debugEnabled = debugEnabled;
        this.benchmarkEnabled = benchmarkEnabled;
        this.lang = normalizeToken(lang, "java");
        this.traceFormat = normalizeTraceFormat(traceFormat);
        this.benchmarkFormat = normalizeOutputFormat(benchmarkFormat);
        this.traceFamilies = Collections.unmodifiableSet(new LinkedHashSet<>(traceFamilies));
        this.traceSignal = traceSignal;
        this.scenario = normalizeToken(scenario, "single_neuron");
        this.value = value;
        this.ticks = Math.max(1, ticks);
        this.outTarget = (outTarget == null || outTarget.isBlank()) ? "stdout" : outTarget;
        this.benchmarkWarmup = Math.max(0, benchmarkWarmup);
        this.benchmarkIterations = Math.max(1, benchmarkIterations);
        this.benchmarkRepeat = Math.max(1, benchmarkRepeat);
        this.benchmarkOutputTarget = (benchmarkOutputTarget == null || benchmarkOutputTarget.isBlank())
                ? "stdout"
                : benchmarkOutputTarget;
        this.traceNeuron = (traceNeuron == null || traceNeuron.isBlank()) ? null : traceNeuron;
        this.traceLayer = traceLayer;
        this.helpRequested = helpRequested;
    }

    public static TraceConfig parse(String lang, String[] args) {
        boolean traceEnabled = false;
        boolean debugEnabled = false;
        boolean benchmarkEnabled = false;
        String traceFormat = "text";
        String benchmarkFormat = "text";
        Set<String> traceFamilies = new LinkedHashSet<>();
        boolean traceSignal = false;
        String scenario = "single_neuron";
        double value = 0.42;
        int ticks = 1;
        String outTarget = "stdout";
        int benchmarkWarmup = 100;
        int benchmarkIterations = 1000;
        int benchmarkRepeat = 5;
        String benchmarkOutputTarget = "stdout";
        String traceNeuron = null;
        Integer traceLayer = null;
        boolean helpRequested = false;

        int argIndex = 0;
        while (argIndex < args.length) {
            String rawArg = args[argIndex];
            if (rawArg == null || rawArg.isBlank()) {
                argIndex += 1;
                continue;
            }
            String optionName = rawArg;
            String optionValue = null;
            int equalsIndex = rawArg.indexOf('=');
            if (equalsIndex >= 0) {
                optionName = rawArg.substring(0, equalsIndex);
                optionValue = rawArg.substring(equalsIndex + 1);
            }

            switch (optionName) {
                case "--lang" -> {
                    String providedLang = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                    String normalizedProvidedLang = normalizeToken(providedLang, lang);
                    if (!normalizeToken(lang, lang).equals(normalizedProvidedLang)) {
                        throw new IllegalArgumentException("This runner only supports --lang " + lang + ".");
                    }
                }
                case "--trace" -> traceEnabled = true;
                case "--debug" -> debugEnabled = true;
                case "--benchmark" -> benchmarkEnabled = true;
                case "--trace-format" -> {
                    traceFormat = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                }
                case "--benchmark-format" -> {
                    benchmarkFormat = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                }
                case "--trace-events" -> {
                    String csvValue = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                    traceFamilies = parseFamilies(csvValue);
                }
                case "--trace-signal" -> traceSignal = true;
                case "--scenario" -> {
                    scenario = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                }
                case "--value" -> {
                    String rawValue = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                    value = Double.parseDouble(rawValue);
                }
                case "--ticks" -> {
                    String rawTicks = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                    ticks = Integer.parseInt(rawTicks);
                }
                case "--benchmark-warmup" -> {
                    String rawWarmup = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                    benchmarkWarmup = Integer.parseInt(rawWarmup);
                }
                case "--benchmark-iterations" -> {
                    String rawIterations = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                    benchmarkIterations = Integer.parseInt(rawIterations);
                }
                case "--benchmark-repeat" -> {
                    String rawRepeat = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                    benchmarkRepeat = Integer.parseInt(rawRepeat);
                }
                case "--out" -> {
                    outTarget = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                }
                case "--benchmark-output" -> {
                    benchmarkOutputTarget = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                }
                case "--trace-neuron" -> {
                    traceNeuron = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                }
                case "--trace-layer" -> {
                    String rawLayer = optionValue != null ? optionValue : nextValue(args, argIndex, optionName);
                    if (optionValue == null) {
                        argIndex += 1;
                    }
                    traceLayer = Integer.parseInt(rawLayer);
                }
                case "--help", "-h" -> helpRequested = true;
                default -> {
                    if (rawArg.startsWith("--")) {
                        throw new IllegalArgumentException("Unknown option: " + rawArg);
                    }
                }
            }
            argIndex += 1;
        }

        return new TraceConfig(
                traceEnabled,
                debugEnabled,
                benchmarkEnabled,
                lang,
                traceFormat,
                benchmarkFormat,
                traceFamilies,
                traceSignal,
                scenario,
                value,
                ticks,
                outTarget,
                benchmarkWarmup,
                benchmarkIterations,
                benchmarkRepeat,
                benchmarkOutputTarget,
                traceNeuron,
                traceLayer,
                helpRequested
        );
    }

    private static String nextValue(String[] args, int argIndex, String optionName) {
        if (argIndex + 1 >= args.length) {
            throw new IllegalArgumentException("Missing value for " + optionName);
        }
        return args[argIndex + 1];
    }

    private static Set<String> parseFamilies(String csvValue) {
        if (csvValue == null || csvValue.isBlank()) {
            return Collections.emptySet();
        }
        Set<String> normalizedFamilies = new LinkedHashSet<>();
        String[] rawParts = csvValue.split(",");
        for (String rawPart : rawParts) {
            String normalizedPart = normalizeToken(rawPart, "");
            if (!normalizedPart.isEmpty()) {
                normalizedFamilies.add(normalizedPart);
            }
        }
        return normalizedFamilies;
    }

    private static String normalizeToken(String rawValue, String fallbackValue) {
        if (rawValue == null || rawValue.isBlank()) {
            return fallbackValue;
        }
        return rawValue.trim().toLowerCase(Locale.ROOT);
    }

    private static String normalizeTraceFormat(String rawFormat) {
        return normalizeOutputFormat(rawFormat);
    }

    private static String normalizeOutputFormat(String rawFormat) {
        String normalizedFormat = normalizeToken(rawFormat, "text");
        return "jsonl".equals(normalizedFormat) ? "jsonl" : "text";
    }

    public static String usage(String lang) {
        String normalizedLang = normalizeToken(lang, "java");
        return String.join(System.lineSeparator(),
                "Usage: grownet trace/benchmark runner",
                "  --lang " + normalizedLang,
                "  --trace",
                "  --benchmark",
                "  --debug",
                "  --trace-format text|jsonl",
                "  --benchmark-format text|jsonl",
                "  --trace-events tick,layer,neuron,slot,weight,threshold,tract,growth,prune,bus",
                "    Exact event-name filters are also accepted, for example: threshold.update",
                "  --trace-signal",
                "  --scenario single_neuron|two_layer|image_2d|growth",
                "  --value 0.42",
                "  --ticks 1",
                "  --benchmark-warmup 100",
                "  --benchmark-iterations 1000",
                "  --benchmark-repeat 5",
                "  --out stdout|path/to/file.log",
                "  --benchmark-output stdout|path/to/file.log",
                "  --trace-neuron E0",
                "  --trace-layer 0",
                "  --help");
    }

    public boolean isTraceEnabled() {
        return traceEnabled;
    }

    public boolean isDebugEnabled() {
        return debugEnabled;
    }

    public boolean isBenchmarkEnabled() {
        return benchmarkEnabled;
    }

    public String getLang() {
        return lang;
    }

    public String getTraceFormat() {
        return traceFormat;
    }

    public String getBenchmarkFormat() {
        return benchmarkFormat;
    }

    public Set<String> getTraceFamilies() {
        return traceFamilies;
    }

    public boolean isTraceSignal() {
        return traceSignal;
    }

    public String getScenario() {
        return scenario;
    }

    public double getValue() {
        return value;
    }

    public int getTicks() {
        return ticks;
    }

    public String getOutTarget() {
        return outTarget;
    }

    public int getBenchmarkWarmup() {
        return benchmarkWarmup;
    }

    public int getBenchmarkIterations() {
        return benchmarkIterations;
    }

    public int getBenchmarkRepeat() {
        return benchmarkRepeat;
    }

    public String getBenchmarkOutputTarget() {
        return benchmarkOutputTarget;
    }

    public boolean writesToStdout() {
        return "stdout".equalsIgnoreCase(outTarget);
    }

    public Path getOutPath() {
        return Path.of(outTarget);
    }

    public boolean writesBenchmarkToStdout() {
        return "stdout".equalsIgnoreCase(benchmarkOutputTarget);
    }

    public Path getBenchmarkOutPath() {
        return Path.of(benchmarkOutputTarget);
    }

    public String getTraceNeuron() {
        return traceNeuron;
    }

    public Integer getTraceLayer() {
        return traceLayer;
    }

    public boolean isHelpRequested() {
        return helpRequested;
    }
}
