package org.ghidra.bridge.query;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import ghidra.program.model.address.Address;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.FunctionIterator;
import ghidra.program.model.listing.Instruction;
import ghidra.program.model.listing.InstructionIterator;
import ghidra.program.model.listing.Program;
import ghidra.program.model.symbol.ExternalLocation;
import ghidra.program.model.symbol.ExternalLocationIterator;
import ghidra.program.model.symbol.Reference;
import ghidra.program.model.symbol.Symbol;
import org.ghidra.bridge.serialization.JsonWriter;

public final class ProgramCallGraph {
    private ProgramCallGraph() {
    }

    public static String json(Program program) {
        return JsonWriter.json(graph(program));
    }

    public static Map<String, Object> graph(Program program) {
        Map<String, Map<String, Object>> nodes = new LinkedHashMap<>();
        Map<String, Map<String, Object>> edges = new LinkedHashMap<>();

        FunctionIterator functions = program.getFunctionManager().getFunctions(true);
        while (functions.hasNext()) {
            Function function = functions.next();
            if (function.isExternal()) {
                continue;
            }
            putNode(nodes, functionNode(function));
        }

        functions = program.getFunctionManager().getFunctions(true);
        while (functions.hasNext()) {
            Function source = functions.next();
            if (source.isExternal()) {
                continue;
            }
            collectEdges(program, source, nodes, edges);
        }

        List<Map<String, Object>> nodeRecords = new ArrayList<>(nodes.values());
        nodeRecords.sort((left, right) -> stringValue(left, "id").compareTo(stringValue(right, "id")));
        List<Map<String, Object>> edgeRecords = new ArrayList<>(edges.values());
        edgeRecords.sort((left, right) -> stringValue(left, "id").compareTo(stringValue(right, "id")));

        Map<String, Object> graph = new LinkedHashMap<>();
        graph.put("nodes", nodeRecords);
        graph.put("edges", edgeRecords);
        return graph;
    }

    private static void collectEdges(
        Program program,
        Function source,
        Map<String, Map<String, Object>> nodes,
        Map<String, Map<String, Object>> edges
    ) {
        Map<String, Object> sourceNode = functionNode(source);
        String sourceId = stringValue(sourceNode, "id");
        InstructionIterator instructions = program.getListing().getInstructions(source.getBody(), true);
        while (instructions.hasNext()) {
            Instruction instruction = instructions.next();
            Reference[] references = program.getReferenceManager().getReferencesFrom(instruction.getAddress());
            for (Reference reference : references) {
                if (!reference.getReferenceType().isCall()) {
                    continue;
                }
                Destination destination = destination(program, reference);
                putNode(nodes, destination.node);
                Map<String, Object> edge = edge(sourceId, destination, instruction.getAddress(), reference);
                edges.putIfAbsent(stringValue(edge, "id"), edge);
            }
        }
    }

    private static Map<String, Object> functionNode(Function function) {
        Map<String, Object> node = new LinkedHashMap<>();
        node.put("id", functionId(function.getEntryPoint()));
        node.put("kind", "function");
        node.put("name", function.getName());
        node.put("address", ProgramModel.address(function.getEntryPoint()));
        node.put("namespace", namespace(function));
        node.put("library", null);
        node.put("function_id", functionId(function.getEntryPoint()));
        node.put("symbol_id", null);
        return node;
    }

    private static Map<String, Object> externalNode(ExternalLocation location, Reference reference) {
        Address address = location.getExternalSpaceAddress();
        if (address == null) {
            address = location.getAddress();
        }
        String name = firstNonEmpty(
            location.getOriginalImportedName(),
            location.getLabel(),
            location.getSymbol() == null ? null : location.getSymbol().getName(),
            reference.getToAddress() == null ? null : reference.getToAddress().toString()
        );
        String library = emptyToNull(location.getLibraryName());

        Map<String, Object> node = new LinkedHashMap<>();
        node.put("id", externalId(library, name, address));
        node.put("kind", "external");
        node.put("name", name);
        node.put("address", ProgramModel.address(address));
        node.put("namespace", null);
        node.put("library", library);
        node.put("function_id", null);
        node.put("symbol_id", null);
        return node;
    }

    private static Map<String, Object> unresolvedNode(Address target) {
        Map<String, Object> node = new LinkedHashMap<>();
        node.put("id", unresolvedId(target));
        node.put("kind", "unresolved");
        node.put("name", target == null ? "unknown" : target.toString());
        node.put("address", ProgramModel.address(target));
        node.put("namespace", null);
        node.put("library", null);
        node.put("function_id", null);
        node.put("symbol_id", null);
        return node;
    }

    private static Map<String, Object> edge(
        String sourceId,
        Destination destination,
        Address callsite,
        Reference reference
    ) {
        Map<String, Object> edge = new LinkedHashMap<>();
        String destinationId = stringValue(destination.node, "id");
        edge.put(
            "id",
            String.join(
                ":",
                "call",
                sourceId,
                callsite.toString(),
                Integer.toString(reference.getOperandIndex()),
                destinationId
            )
        );
        edge.put("source_id", sourceId);
        edge.put("destination_id", destinationId);
        edge.put("callsite", ProgramModel.address(callsite));
        edge.put("reference_type", reference.getReferenceType().getName());
        edge.put("operand_index", reference.getOperandIndex());
        edge.put("computed", computed(reference));
        edge.put("thunk_resolved", destination.thunkResolved);
        return edge;
    }

    private static Destination destination(Program program, Reference reference) {
        Address targetAddress = reference.getToAddress();
        Function target = targetAddress == null
            ? null
            : program.getFunctionManager().getFunctionAt(targetAddress);
        boolean thunkResolved = false;
        if (target != null && target.isThunk()) {
            Function thunked = target.getThunkedFunction(true);
            if (thunked != null) {
                target = thunked;
                thunkResolved = true;
            }
        }
        if (target != null) {
            if (target.isExternal()) {
                ExternalLocation targetLocation = target.getExternalLocation();
                if (targetLocation != null) {
                    return new Destination(externalNode(targetLocation, reference), thunkResolved);
                }
                return new Destination(unresolvedNode(targetAddress), thunkResolved);
            }
            return new Destination(functionNode(target), thunkResolved);
        }

        ExternalLocation location = externalLocationFor(program, reference);
        if (location != null) {
            return new Destination(externalNode(location, reference), false);
        }

        return new Destination(unresolvedNode(targetAddress), false);
    }

    private static ExternalLocation externalLocationFor(Program program, Reference reference) {
        if (reference.isExternalReference()) {
            ExternalLocation location = uniqueExternalLocationAt(program, reference.getToAddress());
            if (location != null) {
                return location;
            }
        }

        Symbol symbol = program.getSymbolTable().getSymbol(reference);
        if (symbol != null && symbol.isExternal()) {
            ExternalLocation location = program.getExternalManager().getExternalLocation(symbol);
            if (location != null) {
                return location;
            }
        }

        Address targetAddress = reference.getToAddress();
        Function target = targetAddress == null ? null : program.getFunctionManager().getFunctionAt(targetAddress);
        if (target == null) {
            return uniqueExternalLocationAt(program, targetAddress);
        }
        if (target.isThunk()) {
            Function thunked = target.getThunkedFunction(true);
            if (thunked != null) {
                target = thunked;
            }
        }
        return target.isExternal() ? target.getExternalLocation() : null;
    }

    private static ExternalLocation uniqueExternalLocationAt(Program program, Address address) {
        if (address == null) {
            return null;
        }
        ExternalLocationIterator locations = program.getExternalManager().getExternalLocations(address);
        if (!locations.hasNext()) {
            return null;
        }
        return locations.next();
    }

    private static void putNode(Map<String, Map<String, Object>> nodes, Map<String, Object> node) {
        nodes.putIfAbsent(stringValue(node, "id"), node);
    }

    private static boolean computed(Reference reference) {
        String name = reference.getReferenceType().getName().toUpperCase(java.util.Locale.ROOT);
        return name.contains("COMPUTED") || name.contains("INDIRECT");
    }

    private static String functionId(Address address) {
        return "function:" + address;
    }

    private static String externalId(String library, String name, Address address) {
        return "external:" + emptyToId(library) + ":" + emptyToId(name) + ":" + emptyToId(address);
    }

    private static String unresolvedId(Address address) {
        return "unresolved:" + emptyToId(address);
    }

    private static String namespace(Function function) {
        return function.getParentNamespace() == null ? null : function.getParentNamespace().getName(true);
    }

    private static String firstNonEmpty(String... values) {
        for (String value : values) {
            value = emptyToNull(value);
            if (value != null) {
                return value;
            }
        }
        return "unknown";
    }

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

    private static String emptyToId(Object value) {
        if (value == null) {
            return "";
        }
        return value.toString();
    }

    private static String stringValue(Map<String, Object> value, String key) {
        return value.get(key).toString();
    }

    private static final class Destination {
        private final Map<String, Object> node;
        private final boolean thunkResolved;

        private Destination(Map<String, Object> node, boolean thunkResolved) {
            this.node = node;
            this.thunkResolved = thunkResolved;
        }
    }
}
