package org.ghidra.bridge.decompile;

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

import ghidra.app.decompiler.DecompInterface;
import ghidra.app.decompiler.DecompiledFunction;
import ghidra.app.decompiler.DecompileResults;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.Program;
import ghidra.program.model.pcode.FunctionPrototype;
import ghidra.program.model.pcode.HighFunction;
import ghidra.program.model.pcode.HighSymbol;
import ghidra.program.model.pcode.LocalSymbolMap;
import ghidra.program.model.pcode.PcodeBlock;
import ghidra.program.model.pcode.PcodeBlockBasic;
import ghidra.program.model.pcode.PcodeOp;
import ghidra.program.model.pcode.Varnode;
import ghidra.util.task.TaskMonitor;
import org.ghidra.bridge.query.ProgramMetadata;
import org.ghidra.bridge.query.ProgramModel;
import org.ghidra.bridge.runtime.TaskMonitors;
import org.ghidra.bridge.serialization.JsonWriter;
import org.ghidra.bridge.serialization.PcodeMaps;
import org.ghidra.bridge.serialization.ProgramTypes;
import org.ghidra.bridge.session.DecompilerHandle;
import org.ghidra.bridge.session.ProgramHandle;

public final class DecompilerService {
    private DecompilerService() {
    }

    public static DecompilerHandle open(ProgramHandle programHandle) {
        DecompInterface decompiler = new DecompInterface();
        decompiler.openProgram(programHandle.program());
        return new DecompilerHandle(programHandle.program(), decompiler);
    }

    public static void close(DecompilerHandle handle) {
        handle.decompiler().dispose();
    }

    public static String decompileJson(
        DecompilerHandle handle,
        String functionAddress,
        int timeout,
        int monitorTimeoutSeconds
    ) {
        Function function = ProgramModel.requireFunctionAt(handle.program(), functionAddress);
        TaskMonitor monitor = TaskMonitors.create(monitorTimeoutSeconds);
        return JsonWriter.json(decompileResultMap(
            handle.decompiler(),
            function,
            timeout,
            monitor,
            handle.types()
        ));
    }

    public static String metadataJson(DecompilerHandle handle) {
        return JsonWriter.json(ProgramMetadata.metadata(handle.program(), handle.types()));
    }

    public static Map<String, Object> decompileResultMap(
        DecompInterface decompiler,
        Function function,
        int timeout,
        TaskMonitor monitor
    ) {
        return decompileResultMap(decompiler, function, timeout, monitor, new ProgramTypes());
    }

    public static Map<String, Object> decompileResultMap(
        DecompInterface decompiler,
        Function function,
        int timeout,
        TaskMonitor monitor,
        ProgramTypes types
    ) {
        DecompileResults result = decompiler.decompileFunction(function, timeout, monitor);
        Map<String, Object> facts = new LinkedHashMap<>();
        String status = decompileStatus(result);
        facts.put("status", status);
        facts.put("error_message", emptyToNull(result.getErrorMessage()));

        if (!"completed".equals(status)) {
            return facts;
        }

        DecompiledFunction decompiled = result.getDecompiledFunction();
        facts.put("c", decompiled == null ? null : decompiled.getC());
        facts.put("signature", decompiled == null ? null : decompiled.getSignature());

        HighFunction high = result.getHighFunction();
        if (high != null) {
            facts.put("prototype", prototype(types, high));
            facts.put("pcode", pcodeSummary(high));
            facts.put("high_pcode", highPcodeGraph(high));
            facts.put("data_flow", DataFlowMaps.graph(high, types));
            facts.put("parameters", parameters(types, high));
            facts.put("local_symbols", localSymbols(types, high));
        }
        return facts;
    }

    private static String decompileStatus(DecompileResults result) {
        if (result.failedToStart()) {
            return "failed_to_start";
        }
        if (result.isTimedOut()) {
            return "timed_out";
        }
        if (result.isCancelled()) {
            return "cancelled";
        }
        if (result.decompileCompleted()) {
            return "completed";
        }
        return "failed";
    }

    private static Map<String, Object> prototype(ProgramTypes types, HighFunction high) {
        FunctionPrototype prototype = high.getFunctionPrototype();
        if (prototype == null) {
            return null;
        }
        Map<String, Object> value = new LinkedHashMap<>();
        value.put(
            "return_type",
            prototype.getReturnType() == null ? "void" : prototype.getReturnType().getDisplayName()
        );
        value.put("return_type_id", types.id(prototype.getReturnType()));
        value.put("calling_convention", emptyToNull(prototype.getModelName()));
        value.put("parameter_count", prototype.getNumParams());
        value.put("varargs", prototype.isVarArg());
        value.put("inline", prototype.isInline());
        value.put("no_return", prototype.hasNoReturn());
        value.put("has_this_pointer", prototype.hasThisPointer());
        return value;
    }

    private static Map<String, Object> pcodeSummary(HighFunction high) {
        Map<String, Long> opCounts = new java.util.TreeMap<>();
        long opCount = 0;
        java.util.Iterator<? extends PcodeOp> ops = high.getPcodeOps();
        while (ops.hasNext()) {
            PcodeOp op = ops.next();
            increment(opCounts, op.getMnemonic());
            opCount++;
        }

        Map<String, Long> varnodeSpaceCounts = new java.util.TreeMap<>();
        long varnodeCount = 0;
        java.util.Iterator<? extends Varnode> varnodes = high.locRange();
        while (varnodes.hasNext()) {
            Varnode varnode = varnodes.next();
            increment(varnodeSpaceCounts, PcodeMaps.varnodeSpace(varnode));
            varnodeCount++;
        }

        Map<String, Object> summary = new LinkedHashMap<>();
        summary.put("op_count", opCount);
        summary.put("op_counts", namedCounts(opCounts));
        summary.put("basic_block_count", high.getBasicBlocks().size());
        summary.put("varnode_count", varnodeCount);
        summary.put("varnode_space_counts", namedCounts(varnodeSpaceCounts));
        return summary;
    }

    private static List<Map<String, Object>> parameters(ProgramTypes types, HighFunction high) {
        LocalSymbolMap symbols = high.getLocalSymbolMap();
        if (symbols == null) {
            return Collections.emptyList();
        }
        List<Map<String, Object>> values = new ArrayList<>();
        for (int i = 0; i < symbols.getNumParams(); i++) {
            Map<String, Object> symbol = highSymbol(types, symbols.getParamSymbol(i));
            if (symbol != null) {
                values.add(symbol);
            }
        }
        return values;
    }

    private static Map<String, Object> highPcodeGraph(HighFunction high) {
        Map<PcodeBlockBasic, String> blockIds = new java.util.IdentityHashMap<>();
        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()));
        });

        List<Map<String, Object>> blockValues = new ArrayList<>();
        for (int i = 0; i < blocks.size(); i++) {
            PcodeBlockBasic block = blocks.get(i);
            String id = "high_block_" + i;
            blockIds.put(block, id);
            blockValues.add(highBlock(block, id));
        }

        List<Map<String, Object>> edgeValues = new ArrayList<>();
        for (PcodeBlockBasic block : blocks) {
            String sourceId = blockIds.get(block);
            for (int i = 0; i < block.getOutSize(); i++) {
                PcodeBlock destination = block.getOut(i);
                String destinationId = destination instanceof PcodeBlockBasic
                    ? blockIds.get((PcodeBlockBasic) destination)
                    : null;
                if (destinationId != null) {
                    edgeValues.add(highEdge(block, sourceId, destination, destinationId));
                }
            }
        }

        Map<String, Object> graph = new LinkedHashMap<>();
        graph.put("blocks", blockValues);
        graph.put("edges", edgeValues);
        graph.put("ops", PcodeMaps.ops(high.getPcodeOps()));
        return graph;
    }

    private static Map<String, Object> highBlock(PcodeBlockBasic block, String id) {
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("id", id);
        value.put("index", block.getIndex());
        value.put("block_type", PcodeBlock.typeToName(block.getType()));
        value.put("range", addressRange(block.getStart(), block.getStop()));
        value.put("op_sequence_numbers", opSequenceNumbers(block));
        return value;
    }

    private static Map<String, Object> highEdge(
        PcodeBlockBasic source,
        String sourceId,
        PcodeBlock destination,
        String destinationId
    ) {
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("source_block_id", sourceId);
        value.put("destination_block_id", destinationId);
        value.put("flow_type", highEdgeType(source, destination));
        value.put("source_address", ProgramModel.address(source.getStop()));
        value.put("destination_address", ProgramModel.address(destination.getStart()));
        value.put("reference_address", null);
        return value;
    }

    private static String highEdgeType(PcodeBlock source, PcodeBlock destination) {
        if (source.getOutSize() >= 2) {
            if (source.getTrueOut() == destination) {
                return "true";
            }
            if (source.getFalseOut() == destination) {
                return "false";
            }
        }
        return "out";
    }

    private static List<String> opSequenceNumbers(PcodeBlockBasic block) {
        List<String> values = new ArrayList<>();
        java.util.Iterator<PcodeOp> ops = block.getIterator();
        while (ops.hasNext()) {
            values.add(ops.next().getSeqnum().toString());
        }
        return values;
    }

    private static Map<String, Object> addressRange(
        ghidra.program.model.address.Address start,
        ghidra.program.model.address.Address end
    ) {
        if (start == null || end == null) {
            return null;
        }
        Map<String, Object> range = new LinkedHashMap<>();
        range.put("start", ProgramModel.address(start));
        range.put("end", ProgramModel.address(end));
        return range;
    }

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

    private static List<Map<String, Object>> localSymbols(ProgramTypes types, HighFunction high) {
        LocalSymbolMap symbols = high.getLocalSymbolMap();
        if (symbols == null) {
            return Collections.emptyList();
        }
        List<Map<String, Object>> values = new ArrayList<>();
        java.util.Iterator<HighSymbol> iterator = symbols.getSymbols();
        while (iterator.hasNext()) {
            HighSymbol symbol = iterator.next();
            if (symbol.isParameter()) {
                continue;
            }
            Map<String, Object> value = highSymbol(types, symbol);
            if (value != null) {
                values.add(value);
            }
        }
        values.sort((left, right) -> symbolSortKey(left).compareTo(symbolSortKey(right)));
        return values;
    }

    private static Map<String, Object> highSymbol(ProgramTypes types, HighSymbol symbol) {
        if (symbol == null || emptyToNull(symbol.getName()) == null) {
            return null;
        }
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("name", symbol.getName());
        value.put(
            "data_type",
            symbol.getDataType() == null ? "undefined" : symbol.getDataType().getDisplayName()
        );
        value.put("type_id", types.id(symbol.getDataType()));
        value.put("size", symbol.getSize());
        value.put("storage", symbol.getStorage() == null ? "unknown" : symbol.getStorage().toString());
        return value;
    }

    private static String symbolSortKey(Map<String, Object> symbol) {
        return symbol.get("name") + "\u001f" + symbol.get("storage") + "\u001f" + symbol.get("data_type");
    }

    private static void increment(Map<String, Long> counts, String name) {
        counts.put(name, counts.getOrDefault(name, 0L) + 1L);
    }

    private static List<Map<String, Object>> namedCounts(Map<String, Long> counts) {
        List<Map<String, Object>> values = new ArrayList<>();
        for (Map.Entry<String, Long> entry : counts.entrySet()) {
            Map<String, Object> value = new LinkedHashMap<>();
            value.put("name", entry.getKey());
            value.put("count", entry.getValue());
            values.add(value);
        }
        return values;
    }

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