package org.ghidra.bridge.decompile;

import java.util.ArrayList;
import java.util.Collections;
import java.util.IdentityHashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

import ghidra.program.model.pcode.HighFunction;
import ghidra.program.model.pcode.HighSymbol;
import ghidra.program.model.pcode.HighVariable;
import ghidra.program.model.pcode.PcodeBlockBasic;
import ghidra.program.model.pcode.PcodeOp;
import ghidra.program.model.pcode.Varnode;
import ghidra.program.model.symbol.Symbol;
import org.ghidra.bridge.query.ProgramModel;
import org.ghidra.bridge.serialization.PcodeMaps;
import org.ghidra.bridge.serialization.ProgramTypes;

public final class DataFlowMaps {
    private DataFlowMaps() {
    }

    public static Map<String, Object> graph(HighFunction high) {
        return graph(high, new ProgramTypes());
    }

    public static Map<String, Object> graph(HighFunction high, ProgramTypes types) {
        List<PcodeBlockBasic> blocks = sortedBlocks(high);
        Map<PcodeBlockBasic, String> blockIds = new IdentityHashMap<>();
        for (int i = 0; i < blocks.size(); i++) {
            blockIds.put(blocks.get(i), "high_block_" + i);
        }

        List<PcodeOp> ops = orderedOps(high, blocks);
        Map<PcodeOp, String> opIds = new IdentityHashMap<>();
        for (int i = 0; i < ops.size(); i++) {
            opIds.put(ops.get(i), "op_" + i);
        }

        Map<String, Varnode> varnodes = new LinkedHashMap<>();
        for (PcodeOp op : ops) {
            collectVarnode(varnodes, op.getOutput());
            for (Varnode input : op.getInputs()) {
                collectVarnode(varnodes, input);
            }
        }

        Map<String, HighVariable> highVariables = new LinkedHashMap<>();
        for (Varnode varnode : varnodes.values()) {
            collectHighVariable(highVariables, varnode);
        }
        for (HighVariable variable : highVariables.values()) {
            for (Varnode instance : variable.getInstances()) {
                collectVarnode(varnodes, instance);
            }
        }

        Map<String, String> varnodeIds = new LinkedHashMap<>();
        int varnodeIndex = 0;
        for (String key : varnodes.keySet()) {
            varnodeIds.put(key, "varnode_" + varnodeIndex);
            varnodeIndex++;
        }

        Map<String, String> variableIds = new LinkedHashMap<>();
        int variableIndex = 0;
        for (String key : highVariables.keySet()) {
            variableIds.put(key, "variable_" + variableIndex);
            variableIndex++;
        }
        for (Varnode varnode : varnodes.values()) {
            String key = syntheticVariableKey(varnode);
            if (variableKey(varnode) == null && !variableIds.containsKey(key)) {
                variableIds.put(key, "variable_" + variableIndex);
                variableIndex++;
            }
        }

        List<Map<String, Object>> opRecords = new ArrayList<>();
        for (PcodeOp op : ops) {
            opRecords.add(opRecord(opIds.get(op), op, blockIds.get(op.getParent())));
        }

        List<Map<String, Object>> variableRecords = new ArrayList<>();
        for (Map.Entry<String, HighVariable> entry : highVariables.entrySet()) {
            variableRecords.add(variableRecord(types, variableIds.get(entry.getKey()), entry.getValue()));
        }
        for (Varnode varnode : varnodes.values()) {
            if (variableKey(varnode) == null) {
                variableRecords.add(syntheticVariableRecord(
                    variableIds.get(syntheticVariableKey(varnode)),
                    varnode
                ));
            }
        }

        List<Map<String, Object>> varnodeRecords = new ArrayList<>();
        for (Map.Entry<String, Varnode> entry : varnodes.entrySet()) {
            varnodeRecords.add(varnodeRecord(
                varnodeIds.get(entry.getKey()),
                entry.getValue(),
                variableIds.get(variableKeyOrSynthetic(entry.getValue()))
            ));
        }

        List<Map<String, Object>> edges = new ArrayList<>();
        for (PcodeOp op : ops) {
            String opId = opIds.get(op);
            for (int i = 0; i < op.getNumInputs(); i++) {
                Varnode input = op.getInput(i);
                if (input != null) {
                    edges.add(edge(varnodeIds.get(varnodeKey(input)), opId, "input", i));
                }
            }
            Varnode output = op.getOutput();
            if (output != null) {
                edges.add(edge(opId, varnodeIds.get(varnodeKey(output)), "output", null));
            }
        }
        for (Varnode varnode : varnodes.values()) {
            String variableId = variableIds.get(variableKeyOrSynthetic(varnode));
            String varnodeId = varnodeIds.get(varnodeKey(varnode));
            if (variableId != null && varnodeId != null) {
                edges.add(edge(variableId, varnodeId, "instance", null));
            }
        }

        Map<String, Object> graph = new LinkedHashMap<>();
        graph.put("ops", opRecords);
        graph.put("varnodes", varnodeRecords);
        graph.put("variables", variableRecords);
        graph.put("edges", edges);
        return graph;
    }

    private static List<PcodeBlockBasic> sortedBlocks(HighFunction high) {
        List<PcodeBlockBasic> blocks = new ArrayList<>();
        for (PcodeBlockBasic block : high.getBasicBlocks()) {
            blocks.add(block);
        }
        blocks.sort((left, right) -> {
            int index = Integer.compare(left.getIndex(), right.getIndex());
            if (index != 0) {
                return index;
            }
            return addressSortKey(left.getStart()).compareTo(addressSortKey(right.getStart()));
        });
        return blocks;
    }

    private static List<PcodeOp> orderedOps(HighFunction high, List<PcodeBlockBasic> blocks) {
        List<PcodeOp> values = new ArrayList<>();
        Set<PcodeOp> seen = Collections.newSetFromMap(new IdentityHashMap<>());
        for (PcodeBlockBasic block : blocks) {
            java.util.Iterator<PcodeOp> iterator = block.getIterator();
            while (iterator.hasNext()) {
                PcodeOp op = iterator.next();
                if (seen.add(op)) {
                    values.add(op);
                }
            }
        }
        java.util.Iterator<? extends PcodeOp> iterator = high.getPcodeOps();
        while (iterator.hasNext()) {
            PcodeOp op = iterator.next();
            if (seen.add(op)) {
                values.add(op);
            }
        }
        return values;
    }

    private static void collectVarnode(Map<String, Varnode> varnodes, Varnode varnode) {
        if (varnode != null) {
            varnodes.putIfAbsent(varnodeKey(varnode), varnode);
        }
    }

    private static void collectHighVariable(Map<String, HighVariable> variables, Varnode varnode) {
        String key = variableKey(varnode);
        if (key != null) {
            variables.putIfAbsent(key, varnode.getHigh());
        }
    }

    private static Map<String, Object> opRecord(String id, PcodeOp op, String blockId) {
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("id", id);
        value.put("sequence_number", op.getSeqnum().toString());
        value.put("mnemonic", op.getMnemonic());
        value.put("opcode", op.getOpcode());
        value.put("block_id", blockId);
        return value;
    }

    private static Map<String, Object> varnodeRecord(String id, Varnode varnode, String variableId) {
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("id", id);
        value.put("kind", varnodeKind(varnode));
        value.put("variable_id", variableId);
        value.put("varnode", PcodeMaps.varnode(varnode));
        return value;
    }

    private static Map<String, Object> variableRecord(ProgramTypes types, String id, HighVariable variable) {
        HighSymbol highSymbol = variable.getSymbol();
        Symbol symbol = highSymbol == null ? null : highSymbol.getSymbol();
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("id", id);
        value.put("name", emptyToFallback(variable.getName(), "unnamed"));
        value.put("data_type", dataTypeName(variable.getDataType()));
        value.put("type_id", types.id(variable.getDataType()));
        value.put("size", variable.getSize());
        value.put("storage", variableStorage(variable));
        value.put("kind", variableKind(variable));
        value.put("symbol_id", highSymbol == null ? null : highSymbol.getId());
        value.put("symbol_address", symbol == null ? null : ProgramModel.address(symbol.getAddress()));
        value.put("symbol_type", symbol == null ? null : symbol.getSymbolType().toString());
        value.put("namespace", highSymbol == null || highSymbol.getNamespace() == null
            ? null
            : highSymbol.getNamespace().getName(true));
        return value;
    }

    private static Map<String, Object> syntheticVariableRecord(String id, Varnode varnode) {
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("id", id);
        value.put("name", varnode.toString());
        value.put("data_type", "undefined");
        value.put("type_id", null);
        value.put("size", varnode.getSize());
        value.put("storage", varnode.toString());
        value.put("kind", varnodeKind(varnode));
        value.put("symbol_id", null);
        value.put("symbol_address", null);
        value.put("symbol_type", null);
        value.put("namespace", null);
        return value;
    }

    private static Map<String, Object> edge(
        String sourceId,
        String destinationId,
        String kind,
        Integer slot
    ) {
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("source_id", sourceId);
        value.put("destination_id", destinationId);
        value.put("kind", kind);
        value.put("slot", slot);
        return value;
    }

    private static String variableStorage(HighVariable variable) {
        HighSymbol symbol = variable.getSymbol();
        if (symbol != null && symbol.getStorage() != null) {
            return symbol.getStorage().toString();
        }
        Varnode representative = variable.getRepresentative();
        return representative == null ? "unknown" : representative.toString();
    }

    private static String variableKind(HighVariable variable) {
        HighSymbol symbol = variable.getSymbol();
        if (symbol != null) {
            if (symbol.isParameter()) {
                return "parameter";
            }
            if (symbol.isGlobal()) {
                return "global";
            }
            return "local";
        }
        Varnode representative = variable.getRepresentative();
        return representative == null ? "unknown" : varnodeKind(representative);
    }

    private static String varnodeKind(Varnode varnode) {
        if (varnode.isConstant()) {
            return "constant";
        }
        if (varnode.isRegister()) {
            return "register";
        }
        if (varnode.isUnique()) {
            return "temporary";
        }
        if (varnode.isPersistent()) {
            return "global";
        }
        if (varnode.isAddress()) {
            return "local";
        }
        return "unknown";
    }

    private static String variableKeyOrSynthetic(Varnode varnode) {
        String key = variableKey(varnode);
        return key == null ? syntheticVariableKey(varnode) : key;
    }

    private static String variableKey(Varnode varnode) {
        HighVariable variable = varnode.getHigh();
        if (variable == null) {
            return null;
        }
        HighSymbol symbol = variable.getSymbol();
        String symbolId = symbol == null ? "" : Long.toString(symbol.getId());
        return "high\u001f" + symbolId
            + "\u001f" + emptyToFallback(variable.getName(), "")
            + "\u001f" + dataTypeName(variable.getDataType())
            + "\u001f" + variable.getSize()
            + "\u001f" + variableStorage(variable)
            + "\u001f" + variableKind(variable);
    }

    private static String syntheticVariableKey(Varnode varnode) {
        return "synthetic\u001f" + varnodeKey(varnode);
    }

    private static String varnodeKey(Varnode varnode) {
        return PcodeMaps.varnodeSpace(varnode)
            + "\u001f" + Long.toUnsignedString(varnode.getOffset())
            + "\u001f" + varnode.getSize()
            + "\u001f" + varnode.toString()
            + "\u001f" + addressSortKey(varnode.getPCAddress())
            + "\u001f" + defKey(varnode)
            + "\u001f" + varnode.getMergeGroup()
            + "\u001f" + emptyToFallback(variableKey(varnode), "");
    }

    private static String dataTypeName(ghidra.program.model.data.DataType dataType) {
        return dataType == null ? "undefined" : dataType.getDisplayName();
    }

    private static String defKey(Varnode varnode) {
        PcodeOp def = varnode.getDef();
        return def == null
            ? ""
            : def.getSeqnum() + "\u001f" + def.getMnemonic() + "\u001f" + def.getOpcode();
    }

    private static String emptyToFallback(String value, String fallback) {
        return value == null || value.isBlank() ? fallback : value;
    }

    private static String addressSortKey(ghidra.program.model.address.Address address) {
        return address == null ? "" : address.toString();
    }
}
