package org.ghidra.bridge.query;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
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.Program;
import ghidra.program.model.mem.MemoryBlock;
import ghidra.program.model.symbol.Reference;
import ghidra.program.model.symbol.ReferenceIterator;
import ghidra.program.model.symbol.Symbol;
import ghidra.program.model.symbol.SymbolIterator;
import org.ghidra.bridge.serialization.JsonWriter;

public final class ProgramModel {
    private ProgramModel() {
    }

    public static String parseAddressJson(Program program, String addressText) {
        return JsonWriter.json(address(parseAddress(program, addressText)));
    }

    public static String functionDescriptorsJson(Program program) {
        List<Map<String, Object>> descriptors = new ArrayList<>();
        FunctionIterator functions = program.getFunctionManager().getFunctions(true);
        while (functions.hasNext()) {
            descriptors.add(functionDescriptor(functions.next()));
        }
        return JsonWriter.json(descriptors);
    }

    public static String functionAtJson(Program program, String addressText) {
        Function function = program.getFunctionManager().getFunctionAt(
            parseAddress(program, addressText)
        );
        return function == null ? null : JsonWriter.json(functionDescriptor(function));
    }

    public static String functionContainingJson(Program program, String addressText) {
        Function function = program.getFunctionManager().getFunctionContaining(
            parseAddress(program, addressText)
        );
        return function == null ? null : JsonWriter.json(functionDescriptor(function));
    }

    public static String symbolsJson(Program program) {
        List<Map<String, Object>> symbols = new ArrayList<>();
        SymbolIterator iterator = program.getSymbolTable().getSymbolIterator();
        while (iterator.hasNext()) {
            Map<String, Object> symbol = symbolMap(iterator.next());
            if (symbol != null) {
                symbols.add(symbol);
            }
        }
        return JsonWriter.json(symbols);
    }

    public static String symbolsAtJson(Program program, String addressText) {
        List<Map<String, Object>> symbols = new ArrayList<>();
        SymbolIterator iterator = program.getSymbolTable().getSymbolsAsIterator(
            parseAddress(program, addressText)
        );
        while (iterator.hasNext()) {
            Map<String, Object> symbol = symbolMap(iterator.next());
            if (symbol != null) {
                symbols.add(symbol);
            }
        }
        return JsonWriter.json(symbols);
    }

    public static String findSymbolsJson(Program program, String name, boolean caseSensitive) {
        List<Map<String, Object>> symbols = new ArrayList<>();
        SymbolIterator iterator = program.getSymbolTable().getSymbolIterator(name, caseSensitive);
        while (iterator.hasNext()) {
            Map<String, Object> symbol = symbolMap(iterator.next());
            if (symbol != null) {
                symbols.add(symbol);
            }
        }
        return JsonWriter.json(symbols);
    }

    public static String referencesFromJson(Program program, String addressText) {
        Reference[] references = program.getReferenceManager().getReferencesFrom(
            parseAddress(program, addressText)
        );
        List<Map<String, Object>> values = new ArrayList<>();
        for (Reference reference : references) {
            values.add(referenceMap(reference));
        }
        return JsonWriter.json(values);
    }

    public static String referencesToJson(Program program, String addressText) {
        ReferenceIterator references = program.getReferenceManager().getReferencesTo(
            parseAddress(program, addressText)
        );
        List<Map<String, Object>> values = new ArrayList<>();
        while (references.hasNext()) {
            values.add(referenceMap(references.next()));
        }
        return JsonWriter.json(values);
    }

    public static String memoryBlocksJson(Program program) {
        List<Map<String, Object>> values = new ArrayList<>();
        for (MemoryBlock block : program.getMemory().getBlocks()) {
            values.add(memoryBlockMap(block));
        }
        return JsonWriter.json(values);
    }

    public static Map<String, Object> functionDescriptor(Function function) {
        Map<String, Object> descriptor = new LinkedHashMap<>();
        descriptor.put("entry", address(function.getEntryPoint()));
        descriptor.put("name", function.getName());
        return descriptor;
    }

    public static Map<String, Object> address(Address address) {
        if (address == null) {
            return null;
        }
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("space", address.getAddressSpace().getName());
        value.put("offset", Long.toUnsignedString(address.getOffset()));
        value.put("text", address.toString());
        return value;
    }

    public static Address parseAddress(Program program, String addressText) {
        Address address = program.getAddressFactory().getAddress(addressText);
        if (address == null) {
            throw new IllegalArgumentException("invalid address " + addressText);
        }
        return address;
    }

    public static Function requireFunctionAt(Program program, String addressText) {
        Address address = parseAddress(program, addressText);
        Function function = program.getFunctionManager().getFunctionAt(address);
        if (function == null) {
            throw new IllegalArgumentException("no function at address " + addressText);
        }
        return function;
    }

    private static Map<String, Object> symbolMap(Symbol symbol) {
        if (symbol.getAddress() == null) {
            return null;
        }
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("name", symbol.getName());
        value.put("address", address(symbol.getAddress()));
        value.put("kind", symbolKind(symbol.getSymbolType().toString()));
        value.put("source", symbolSource(symbol.getSource().toString()));
        value.put(
            "namespace",
            symbol.getParentNamespace() == null ? "" : symbol.getParentNamespace().toString()
        );
        value.put("primary", symbol.isPrimary());
        value.put("external", symbol.isExternal());
        return value;
    }

    public static Map<String, Object> referenceMap(Reference reference) {
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("from", address(reference.getFromAddress()));
        value.put("to", address(reference.getToAddress()));
        value.put("reference_type", reference.getReferenceType().getName());
        value.put("operand_index", reference.getOperandIndex());
        value.put("primary", reference.isPrimary());
        value.put("external", reference.isExternalReference());
        value.put("call", reference.getReferenceType().isCall());
        value.put("data", reference.getReferenceType().isData());
        value.put("read", reference.getReferenceType().isRead());
        value.put("write", reference.getReferenceType().isWrite());
        return value;
    }

    private static Map<String, Object> memoryBlockMap(MemoryBlock block) {
        Map<String, Object> value = new LinkedHashMap<>();
        Map<String, Object> range = new LinkedHashMap<>();
        range.put("start", address(block.getStart()));
        range.put("end", address(block.getEnd()));
        value.put("name", block.getName());
        value.put("range", range);
        value.put("size", block.getSize());
        value.put("initialized", block.isInitialized());
        value.put("read", block.isRead());
        value.put("write", block.isWrite());
        value.put("execute", block.isExecute());
        value.put("volatile", block.isVolatile());
        value.put("overlay", block.isOverlay());
        return value;
    }

    private static String symbolKind(String kind) {
        switch (normalized(kind)) {
            case "FUNCTION":
                return "function";
            case "LABEL":
                return "label";
            case "NAMESPACE":
                return "namespace";
            case "CLASS":
                return "class";
            case "LIBRARY":
                return "library";
            case "PARAMETER":
                return "parameter";
            case "LOCALVAR":
            case "LOCALVARIABLE":
                return "local_variable";
            case "GLOBAL":
            case "GLOBALVAR":
            case "GLOBALVARIABLE":
                return "global_variable";
            default:
                return "unknown";
        }
    }

    private static String symbolSource(String source) {
        switch (normalized(source)) {
            case "USERDEFINED":
                return "user_defined";
            case "IMPORTED":
                return "imported";
            case "ANALYSIS":
                return "analysis";
            case "DEFAULT":
                return "default";
            default:
                return "unknown";
        }
    }

    private static String normalized(String value) {
        return value == null
            ? ""
            : value.replace("_", "").replace(" ", "").replace("-", "").toUpperCase(Locale.ROOT);
    }
}
