package org.ghidra.bridge.query;

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

import ghidra.program.model.address.Address;
import ghidra.program.model.address.AddressSet;
import ghidra.program.model.block.BasicBlockModel;
import ghidra.program.model.block.CodeBlock;
import ghidra.program.model.block.CodeBlockIterator;
import ghidra.program.model.block.CodeBlockReference;
import ghidra.program.model.block.CodeBlockReferenceIterator;
import ghidra.program.model.listing.Data;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.Instruction;
import ghidra.program.model.listing.InstructionIterator;
import ghidra.program.model.listing.Program;
import ghidra.program.model.mem.MemoryAccessException;
import ghidra.program.model.scalar.Scalar;
import ghidra.program.model.symbol.Reference;
import ghidra.util.task.TaskMonitor;
import org.ghidra.bridge.serialization.JsonWriter;
import org.ghidra.bridge.serialization.PcodeMaps;

public final class ProgramListing {
    private ProgramListing() {
    }

    public static String instructionsForFunctionJson(Program program, String functionAddress) throws Exception {
        Function function = ProgramModel.requireFunctionAt(program, functionAddress);
        return JsonWriter.json(instructionsForFunction(program, function));
    }

    public static String instructionsInRangeJson(Program program, String startText, String endText) throws Exception {
        Address start = ProgramModel.parseAddress(program, startText);
        Address end = ProgramModel.parseAddress(program, endText);
        if (!start.getAddressSpace().equals(end.getAddressSpace())) {
            throw new IllegalArgumentException("instruction range must stay within one address space");
        }
        if (start.compareTo(end) > 0) {
            throw new IllegalArgumentException("instruction range start must not be after end");
        }
        List<Map<String, Object>> values = new ArrayList<>();
        InstructionIterator instructions = program.getListing().getInstructions(
            new AddressSet(start, end),
            true
        );
        while (instructions.hasNext()) {
            values.add(instructionMap(program, instructions.next()));
        }
        return JsonWriter.json(values);
    }

    public static String instructionAtJson(Program program, String addressText) throws Exception {
        Instruction instruction = program.getListing().getInstructionAt(
            ProgramModel.parseAddress(program, addressText)
        );
        return instruction == null ? null : JsonWriter.json(instructionMap(program, instruction));
    }

    public static String dataAtJson(Program program, String addressText) {
        Data data = program.getListing().getDataAt(ProgramModel.parseAddress(program, addressText));
        return data == null ? null : JsonWriter.json(dataMap(data));
    }

    public static String dataContainingJson(Program program, String addressText) {
        Data data = program.getListing().getDataContaining(ProgramModel.parseAddress(program, addressText));
        return data == null ? null : JsonWriter.json(dataMap(data));
    }

    public static String controlFlowGraphJson(
        Program program,
        String functionAddress,
        TaskMonitor monitor
    ) throws Exception {
        Function function = ProgramModel.requireFunctionAt(program, functionAddress);
        return JsonWriter.json(controlFlowGraph(program, new BasicBlockModel(program), function, monitor));
    }

    public static String controlFlowGraphsJson(
        Program program,
        String[] functionAddresses,
        TaskMonitor monitor
    ) throws Exception {
        BasicBlockModel model = new BasicBlockModel(program);
        List<Map<String, Object>> values = new ArrayList<>();
        for (String functionAddress : functionAddresses) {
            Function function = ProgramModel.requireFunctionAt(program, functionAddress);
            Map<String, Object> value = new LinkedHashMap<>();
            value.put("function", ProgramModel.functionDescriptor(function));
            value.put("graph", controlFlowGraph(program, model, function, monitor));
            values.add(value);
        }
        return JsonWriter.json(values);
    }

    public static Map<String, Object> controlFlowGraph(
        Program program,
        BasicBlockModel model,
        Function function,
        TaskMonitor monitor
    ) throws Exception {
        List<CodeBlock> blocks = new ArrayList<>();
        CodeBlockIterator iterator = model.getCodeBlocksContaining(function.getBody(), monitor);
        while (iterator.hasNext()) {
            blocks.add(iterator.next());
        }
        blocks.sort((left, right) -> left.getFirstStartAddress().compareTo(right.getFirstStartAddress()));

        Map<String, String> blockIds = new java.util.HashMap<>();
        List<Map<String, Object>> blockValues = new ArrayList<>();
        for (int i = 0; i < blocks.size(); i++) {
            CodeBlock block = blocks.get(i);
            String id = "block_" + i;
            blockIds.put(blockKey(block), id);
            blockValues.add(basicBlock(program, block, id));
        }

        List<Map<String, Object>> edgeValues = new ArrayList<>();
        for (CodeBlock block : blocks) {
            String sourceId = blockIds.get(blockKey(block));
            CodeBlockReferenceIterator destinations = block.getDestinations(monitor);
            while (destinations.hasNext()) {
                CodeBlockReference reference = destinations.next();
                CodeBlock destination = reference.getDestinationBlock();
                String destinationId = destination == null ? null : blockIds.get(blockKey(destination));
                if (destinationId != null) {
                    edgeValues.add(edge(sourceId, destinationId, reference));
                }
            }
        }

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

    public static long basicBlockCount(
        BasicBlockModel model,
        Function function,
        TaskMonitor monitor
    ) throws Exception {
        CodeBlockIterator blocks = model.getCodeBlocksContaining(function.getBody(), monitor);
        long count = 0;
        while (blocks.hasNext()) {
            blocks.next();
            count++;
        }
        return count;
    }

    public static long instructionCount(Program program, Function function) {
        InstructionIterator instructions = program.getListing().getInstructions(function.getBody(), true);
        long count = 0;
        while (instructions.hasNext()) {
            instructions.next();
            count++;
        }
        return count;
    }

    public static List<String> strings(Program program, Function function) {
        Set<String> strings = new TreeSet<>();
        for (Reference reference : nonCallReferencesFrom(program, function)) {
            String value = stringAt(program, reference.getToAddress());
            if (value != null && !value.isEmpty()) {
                strings.add(value);
            }
        }
        return new ArrayList<>(strings);
    }

    public static List<String> constants(Program program, Function function) {
        Set<String> constants = new TreeSet<>();
        InstructionIterator instructions = program.getListing().getInstructions(function.getBody(), true);
        while (instructions.hasNext()) {
            Instruction instruction = instructions.next();
            for (int i = 0; i < instruction.getNumOperands(); i++) {
                for (Object value : instruction.getOpObjects(i)) {
                    Long scalar = scalarValue(value);
                    if (scalar != null) {
                        constants.add(hex(scalar));
                    }
                }
            }
        }
        return new ArrayList<>(constants);
    }

    public static List<Map<String, Object>> dataRefs(Program program, Function function) {
        Map<String, Address> refs = new java.util.TreeMap<>();
        for (Reference reference : nonCallReferencesFrom(program, function)) {
            refs.put(reference.getToAddress().toString(), reference.getToAddress());
        }
        return addresses(refs);
    }

    public static List<Reference> nonCallReferencesFrom(Program program, Function function) {
        List<Reference> refs = new ArrayList<>();
        InstructionIterator instructions = program.getListing().getInstructions(function.getBody(), true);
        while (instructions.hasNext()) {
            Instruction instruction = instructions.next();
            Reference[] references = program.getReferenceManager().getReferencesFrom(instruction.getAddress());
            for (Reference reference : references) {
                if (!reference.getReferenceType().isCall()) {
                    refs.add(reference);
                }
            }
        }
        return refs;
    }

    private static List<Map<String, Object>> instructionsForFunction(
        Program program,
        Function function
    ) throws Exception {
        List<Map<String, Object>> values = new ArrayList<>();
        InstructionIterator instructions = program.getListing().getInstructions(function.getBody(), true);
        while (instructions.hasNext()) {
            values.add(instructionMap(program, instructions.next()));
        }
        return values;
    }

    private static Map<String, Object> instructionMap(
        Program program,
        Instruction instruction
    ) throws MemoryAccessException {
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("address", ProgramModel.address(instruction.getAddress()));
        value.put("mnemonic", instruction.getMnemonicString());
        value.put("pcode", PcodeMaps.ops(instruction.getPcode()));
        value.put("operands", operands(program, instruction));
        value.put("length", instruction.getLength());
        value.put("bytes", bytes(instruction.getBytes()));
        value.put("fallthrough", ProgramModel.address(instruction.getFallThrough()));
        value.put("flow_type", instruction.getFlowType().toString());
        return value;
    }

    private static Map<String, Object> basicBlock(
        Program program,
        CodeBlock block,
        String id
    ) {
        Map<String, Object> value = new LinkedHashMap<>();
        Map<String, Object> range = new LinkedHashMap<>();
        range.put("start", ProgramModel.address(block.getMinAddress()));
        range.put("end", ProgramModel.address(block.getMaxAddress()));
        value.put("id", id);
        value.put("name", block.getName());
        value.put("range", range);
        value.put("flow_type", block.getFlowType().toString());
        value.put("instruction_addresses", instructionAddresses(program, block));
        return value;
    }

    private static List<Map<String, Object>> instructionAddresses(Program program, CodeBlock block) {
        List<Map<String, Object>> values = new ArrayList<>();
        InstructionIterator instructions = program.getListing().getInstructions(block, true);
        while (instructions.hasNext()) {
            values.add(ProgramModel.address(instructions.next().getAddress()));
        }
        return values;
    }

    private static Map<String, Object> edge(
        String sourceId,
        String destinationId,
        CodeBlockReference reference
    ) {
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("source_block_id", sourceId);
        value.put("destination_block_id", destinationId);
        value.put("flow_type", reference.getFlowType().toString());
        value.put("source_address", ProgramModel.address(reference.getSourceAddress()));
        value.put("destination_address", ProgramModel.address(reference.getDestinationAddress()));
        value.put("reference_address", ProgramModel.address(reference.getReference()));
        return value;
    }

    private static String blockKey(CodeBlock block) {
        return block.getFirstStartAddress().toString();
    }

    private static List<Map<String, Object>> operands(Program program, Instruction instruction) {
        List<Map<String, Object>> values = new ArrayList<>();
        for (int i = 0; i < instruction.getNumOperands(); i++) {
            Map<String, Object> operand = new LinkedHashMap<>();
            operand.put("index", i);
            operand.put("text", instruction.getDefaultOperandRepresentation(i));
            operand.put("objects", operandObjects(instruction.getOpObjects(i)));
            operand.put("references", operandReferences(program, instruction, i));
            values.add(operand);
        }
        return values;
    }

    private static List<String> operandObjects(Object[] objects) {
        List<String> values = new ArrayList<>();
        for (Object object : objects) {
            if (object != null) {
                values.add(object.toString());
            }
        }
        return values;
    }

    private static List<Map<String, Object>> operandReferences(
        Program program,
        Instruction instruction,
        int operandIndex
    ) {
        List<Map<String, Object>> values = new ArrayList<>();
        Reference[] references = program.getReferenceManager().getReferencesFrom(instruction.getAddress());
        for (Reference reference : references) {
            if (reference.getOperandIndex() == operandIndex) {
                values.add(ProgramModel.referenceMap(reference));
            }
        }
        return values;
    }

    private static Map<String, Object> dataMap(Data data) {
        Map<String, Object> value = new LinkedHashMap<>();
        Map<String, Object> range = new LinkedHashMap<>();
        range.put("start", ProgramModel.address(data.getMinAddress()));
        range.put("end", ProgramModel.address(data.getMaxAddress()));
        Object dataValue = data.getValue();
        value.put("address", ProgramModel.address(data.getAddress()));
        value.put("range", range);
        value.put("data_type", data.getDataType().getDisplayName());
        value.put("length", data.getLength());
        value.put("value", dataValue == null ? null : dataValue.toString());
        value.put("display", data.getDefaultValueRepresentation());
        return value;
    }

    private static String stringAt(Program program, Address address) {
        Data data = program.getListing().getDataAt(address);
        if (data == null) {
            data = program.getListing().getDataContaining(address);
        }
        if (data == null || data.getValue() == null) {
            return null;
        }

        Object value = data.getValue();
        String dataType = data.getDataType().getName().toLowerCase();
        if (!(value instanceof String) && !dataType.contains("string")) {
            return null;
        }
        return value.toString();
    }

    private static Long scalarValue(Object value) {
        if (value instanceof Number) {
            return ((Number) value).longValue();
        }
        if (value instanceof Scalar) {
            return ((Scalar) value).getValue();
        }
        return null;
    }

    private static String hex(long value) {
        if (value < 0) {
            return "-0x" + Long.toHexString(-value);
        }
        return "0x" + Long.toHexString(value);
    }

    private static List<Map<String, Object>> addresses(Map<String, Address> sortedAddresses) {
        List<Map<String, Object>> values = new ArrayList<>();
        for (Address address : sortedAddresses.values()) {
            values.add(ProgramModel.address(address));
        }
        return values;
    }

    private static List<Integer> bytes(byte[] bytes) {
        List<Integer> values = new ArrayList<>();
        for (byte value : bytes) {
            values.add(Byte.toUnsignedInt(value));
        }
        return values;
    }
}
