package ai.nektron.grownet.debug;

import ai.nektron.grownet.Layer;
import ai.nektron.grownet.Neuron;
import ai.nektron.grownet.Region;
import ai.nektron.grownet.metrics.RegionMetrics;

import java.io.PrintWriter;
import java.io.Writer;
import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;

public final class Trace {
    private static final List<String> FIELD_ORDER = List.of(
            "tick",
            "lang",
            "scenario",
            "region",
            "layer_index",
            "layer_name",
            "source_layer_index",
            "dest_layer_index",
            "neuron_id",
            "source_neuron_id",
            "target_neuron_id",
            "slot_id",
            "signal_id",
            "input_port",
            "input_value",
            "output_value",
            "fallback",
            "created_slot",
            "created_neuron",
            "created_layer",
            "reason",
            "hop",
            "desired_slot",
            "chosen_slot",
            "slots_before",
            "slots_after",
            "slot_limit",
            "neuron_count",
            "strength_before",
            "strength_after",
            "modulation",
            "hit_count",
            "theta_before",
            "theta_after",
            "fired",
            "delivered_events",
            "total_slots",
            "total_synapses",
            "step"
    );

    private static final ThreadLocal<TraceState> STATE = ThreadLocal.withInitial(TraceState::new);
    private static TraceConfig config = TraceConfig.parse("java", new String[0]);
    private static PrintWriter writer = new PrintWriter(System.out, true);

    private Trace() {
    }

    public static void configure(TraceConfig traceConfig, Writer traceWriter) {
        config = Objects.requireNonNull(traceConfig, "traceConfig");
        writer = new PrintWriter(Objects.requireNonNull(traceWriter, "traceWriter"), true);
        STATE.set(new TraceState());
    }

    public static void reset() {
        STATE.remove();
    }

    public static TraceConfig getConfig() {
        return config;
    }

    public static boolean isEnabled() {
        return config != null && config.isTraceEnabled();
    }

    public static int beginTick(String regionName, String inputPortName, Double inputValue, boolean useDefaultSignalId) {
        TraceState traceState = STATE.get();
        traceState.tickNumber += 1;
        traceState.regionName = regionName;
        traceState.signalId = useDefaultSignalId ? formatTickSignalId(traceState.tickNumber) : null;
        traceState.hop = 0;
        Map<String, Object> eventFields = new LinkedHashMap<>();
        if (inputPortName != null) {
            eventFields.put("input_port", inputPortName);
        }
        if (inputValue != null) {
            eventFields.put("input_value", inputValue);
        }
        emit("tick.start", eventFields);
        return traceState.tickNumber;
    }

    public static void endTick(RegionMetrics metrics) {
        if (metrics == null) {
            STATE.get().signalId = null;
            STATE.get().hop = 0;
            return;
        }
        Map<String, Object> eventFields = new LinkedHashMap<>();
        eventFields.put("delivered_events", metrics.getDeliveredEvents());
        eventFields.put("total_slots", metrics.getTotalSlots());
        eventFields.put("total_synapses", metrics.getTotalSynapses());
        emit("tick.end", eventFields);
        TraceState traceState = STATE.get();
        traceState.signalId = null;
        traceState.hop = 0;
    }

    public static void runWithSignal(String signalId, Runnable runnable) {
        TraceState traceState = STATE.get();
        String previousSignalId = traceState.signalId;
        int previousHop = traceState.hop;
        traceState.signalId = signalId;
        traceState.hop = 0;
        try {
            runnable.run();
        } finally {
            traceState.signalId = previousSignalId;
            traceState.hop = previousHop;
        }
    }

    public static void runWithIncrementedHop(Runnable runnable) {
        TraceState traceState = STATE.get();
        int previousHop = traceState.hop;
        traceState.hop = previousHop + 1;
        try {
            runnable.run();
        } finally {
            traceState.hop = previousHop;
        }
    }

    public static String currentSignalId() {
        return STATE.get().signalId;
    }

    public static int currentTick() {
        return STATE.get().tickNumber;
    }

    public static int currentHop() {
        return STATE.get().hop;
    }

    public static String formatTickSignalId(int tickNumber) {
        return String.format(Locale.ROOT, "tick-%06d", tickNumber);
    }

    public static String formatPixelSignalId(int tickNumber, int rowIndex, int colIndex) {
        return String.format(Locale.ROOT, "tick-%06d:px-%04d-%04d", tickNumber, rowIndex, colIndex);
    }

    public static void addLayerFields(Map<String, Object> eventFields, Layer layer) {
        if (eventFields == null || layer == null) {
            return;
        }
        Region region = resolveRegion(layer);
        if (region == null) {
            return;
        }
        int layerIndex = region.getLayers().indexOf(layer);
        if (layerIndex >= 0) {
            eventFields.put("layer_index", layerIndex);
            eventFields.put("layer_name", "layer_" + layerIndex);
        }
    }

    public static void addNeuronFields(Map<String, Object> eventFields, Neuron neuron) {
        if (eventFields == null || neuron == null) {
            return;
        }
        eventFields.put("neuron_id", neuron.getId());
        addLayerFields(eventFields, neuron.owner);
    }

    public static void emitBusDecay(Layer layer) {
        if (layer == null) {
            return;
        }
        Map<String, Object> eventFields = new LinkedHashMap<>();
        addLayerFields(eventFields, layer);
        eventFields.put("step", layer.getBus().getCurrentStep());
        emit("bus.decay", eventFields);
    }

    public static void emitLayerForward(Layer layer, double inputValue) {
        if (layer == null) {
            return;
        }
        Map<String, Object> eventFields = new LinkedHashMap<>();
        addLayerFields(eventFields, layer);
        eventFields.put("neuron_count", layer.getNeurons().size());
        eventFields.put("input_value", inputValue);
        emit("layer.forward", eventFields);
    }

    public static void emit(String eventName, Map<String, Object> rawFields) {
        if (!isEnabled()) {
            return;
        }
        TraceState traceState = STATE.get();
        if (traceState.tickNumber <= 0) {
            return;
        }

        LinkedHashMap<String, Object> eventFields = new LinkedHashMap<>();
        eventFields.put("event", eventName);
        eventFields.put("tick", traceState.tickNumber);
        eventFields.put("lang", config.getLang());
        eventFields.put("scenario", config.getScenario());
        eventFields.put("region", traceState.regionName);

        if (rawFields != null) {
            for (Map.Entry<String, Object> fieldEntry : rawFields.entrySet()) {
                if (fieldEntry.getValue() != null) {
                    eventFields.put(fieldEntry.getKey(), fieldEntry.getValue());
                }
            }
        }

        if (!eventFields.containsKey("signal_id") && traceState.signalId != null) {
            eventFields.put("signal_id", traceState.signalId);
        }
        if (!eventFields.containsKey("hop") && traceState.signalId != null) {
            eventFields.put("hop", traceState.hop);
        }

        if (!passesFilters(eventName, eventFields)) {
            return;
        }

        try {
            if ("jsonl".equals(config.getTraceFormat())) {
                writer.println(formatJsonLine(eventFields));
            } else {
                writer.println(formatTextLine(eventFields));
            }
            writer.flush();
        } catch (Throwable ignored) {
        }
    }

    private static boolean passesFilters(String eventName, Map<String, Object> eventFields) {
        if (config == null) {
            return false;
        }
        if (!config.getTraceFamilies().isEmpty()) {
            if (!matchesEventFilter(eventName, config.getTraceFamilies())) {
                return false;
            }
        }
        String filteredNeuronId = config.getTraceNeuron();
        if (filteredNeuronId != null) {
            Object neuronId = firstPresent(eventFields, List.of("neuron_id", "source_neuron_id", "target_neuron_id"));
            if (neuronId == null || !filteredNeuronId.equals(String.valueOf(neuronId))) {
                return false;
            }
        }
        Integer filteredLayerIndex = config.getTraceLayer();
        if (filteredLayerIndex != null) {
            Object layerIndex = firstPresent(eventFields, List.of("layer_index", "source_layer_index", "dest_layer_index"));
            if (layerIndex == null || filteredLayerIndex.intValue() != parseInt(String.valueOf(layerIndex), Integer.MIN_VALUE)) {
                return false;
            }
        }
        return true;
    }

    private static Region resolveRegion(Layer layer) {
        try {
            Field regionField = Layer.class.getDeclaredField("region");
            regionField.setAccessible(true);
            Object rawRegion = regionField.get(layer);
            if (rawRegion instanceof Region region) {
                return region;
            }
        } catch (Throwable ignored) {
        }
        return null;
    }

    private static boolean matchesEventFilter(String eventName, Collection<String> filterTokens) {
        String normalizedEventName = eventName == null ? "" : eventName.toLowerCase(Locale.ROOT);
        String familyName = normalizedEventName;
        int separatorIndex = normalizedEventName.indexOf('.');
        if (separatorIndex > 0) {
            familyName = normalizedEventName.substring(0, separatorIndex);
        }
        for (String filterToken : filterTokens) {
            if (filterToken == null) {
                continue;
            }
            String normalizedFilterToken = filterToken.toLowerCase(Locale.ROOT);
            if (normalizedEventName.equals(normalizedFilterToken) || familyName.equals(normalizedFilterToken)) {
                return true;
            }
        }
        return false;
    }

    private static Object firstPresent(Map<String, Object> eventFields, Collection<String> fieldNames) {
        for (String fieldName : fieldNames) {
            Object fieldValue = eventFields.get(fieldName);
            if (fieldValue != null) {
                return fieldValue;
            }
        }
        return null;
    }

    private static int parseInt(String rawValue, int fallbackValue) {
        try {
            return Integer.parseInt(rawValue);
        } catch (NumberFormatException ignored) {
            return fallbackValue;
        }
    }

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

    private static String formatJsonLine(Map<String, Object> eventFields) {
        StringBuilder builder = new StringBuilder();
        builder.append('{');
        boolean firstField = true;
        for (String fieldName : orderedKeys(eventFields)) {
            Object fieldValue = eventFields.get(fieldName);
            if (fieldValue == null) {
                continue;
            }
            if (!firstField) {
                builder.append(',');
            }
            firstField = false;
            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) && !"event".equals(fieldName)) {
                additionalKeys.add(fieldName);
            }
        }
        additionalKeys.sort(Comparator.naturalOrder());
        orderedKeys.addAll(additionalKeys);
        orderedKeys.add(0, "event");
        return orderedKeys;
    }

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

    private static void appendJsonValue(StringBuilder builder, Object value) {
        if (value instanceof Number || value instanceof Boolean) {
            builder.append(String.valueOf(value));
            return;
        }
        builder.append('"').append(escapeJson(String.valueOf(value))).append('"');
    }

    private static String escapeJson(String rawValue) {
        StringBuilder builder = new StringBuilder();
        for (int index = 0; index < rawValue.length(); index += 1) {
            char currentChar = rawValue.charAt(index);
            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 final class TraceState {
        private int tickNumber = 0;
        private String regionName = null;
        private String signalId = null;
        private int hop = 0;
    }
}
