package org.ghidra.bridge.query;

import java.util.ArrayList;
import java.util.Collections;
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.block.BasicBlockModel;
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.symbol.ExternalLocation;
import ghidra.program.model.symbol.ExternalLocationIterator;
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 ghidra.util.task.TaskMonitor;
import org.ghidra.bridge.runtime.TaskMonitors;
import org.ghidra.bridge.serialization.JsonWriter;

public final class ProgramFunctionRelations {
    private ProgramFunctionRelations() {
    }

    public static String callersJson(Program program, String functionAddress) {
        return JsonWriter.json(callers(program, ProgramModel.requireFunctionAt(program, functionAddress)));
    }

    public static String calleesJson(Program program, String functionAddress) {
        return JsonWriter.json(callees(program, ProgramModel.requireFunctionAt(program, functionAddress)));
    }

    public static String stringsJson(Program program, String functionAddress) {
        Function function = ProgramModel.requireFunctionAt(program, functionAddress);
        return JsonWriter.json(ProgramListing.strings(program, function));
    }

    public static String constantsJson(Program program, String functionAddress) {
        Function function = ProgramModel.requireFunctionAt(program, functionAddress);
        return JsonWriter.json(ProgramListing.constants(program, function));
    }

    public static String dataRefsJson(Program program, String functionAddress) {
        Function function = ProgramModel.requireFunctionAt(program, functionAddress);
        return JsonWriter.json(ProgramListing.dataRefs(program, function));
    }

    public static String importsJson(Program program, String functionAddress) {
        return JsonWriter.json(imports(program, ProgramModel.requireFunctionAt(program, functionAddress)));
    }

    public static String exportsJson(Program program, String functionAddress) {
        return JsonWriter.json(exports(program, ProgramModel.requireFunctionAt(program, functionAddress)));
    }

    public static String summaryJson(Program program, String functionAddress) throws Exception {
        Function function = ProgramModel.requireFunctionAt(program, functionAddress);
        return JsonWriter.json(summary(program, new BasicBlockModel(program), function, TaskMonitors.create(-1)));
    }

    public static String summariesJson(Program program, String[] functionAddresses) throws Exception {
        BasicBlockModel basicBlockModel = new BasicBlockModel(program);
        TaskMonitor monitor = TaskMonitors.create(-1);
        List<Map<String, Object>> values = new ArrayList<>();
        for (String functionAddress : functionAddresses) {
            values.add(summary(
                program,
                basicBlockModel,
                ProgramModel.requireFunctionAt(program, functionAddress),
                monitor
            ));
        }
        return JsonWriter.json(values);
    }

    private static Map<String, Object> summary(
        Program program,
        BasicBlockModel basicBlockModel,
        Function function,
        TaskMonitor monitor
    ) throws Exception {
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("function", ProgramModel.functionDescriptor(function));
        value.put("basic_block_count", ProgramListing.basicBlockCount(basicBlockModel, function, monitor));
        value.put("instruction_count", ProgramListing.instructionCount(program, function));
        value.put("callers", callers(program, function));
        value.put("callees", callees(program, function));
        value.put("strings", ProgramListing.strings(program, function));
        value.put("constants", ProgramListing.constants(program, function));
        value.put("data_refs", ProgramListing.dataRefs(program, function));
        value.put("imports", imports(program, function));
        value.put("exports", exports(program, function));
        value.put("source_map", ProgramSourceMaps.sourceMap(program, function));
        return value;
    }

    private static List<Map<String, Object>> callers(Program program, Function function) {
        Map<String, Address> callers = new java.util.TreeMap<>();
        ReferenceIterator references = program.getReferenceManager().getReferencesTo(function.getEntryPoint());
        while (references.hasNext()) {
            Reference reference = references.next();
            Function source = program.getFunctionManager().getFunctionContaining(reference.getFromAddress());
            if (source != null) {
                callers.put(source.getEntryPoint().toString(), source.getEntryPoint());
            }
        }
        return addresses(callers);
    }

    private static List<Map<String, Object>> callees(Program program, Function function) {
        Map<String, Address> callees = new java.util.TreeMap<>();
        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()) {
                    continue;
                }
                Function target = program.getFunctionManager().getFunctionAt(reference.getToAddress());
                if (target != null) {
                    callees.put(target.getEntryPoint().toString(), target.getEntryPoint());
                }
            }
        }
        return addresses(callees);
    }

    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<Map<String, Object>> imports(Program program, Function function) {
        Map<String, ImportRecord> imports = new LinkedHashMap<>();
        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) {
                ExternalLocation location = externalLocationFor(program, reference);
                if (location == null) {
                    continue;
                }
                ImportRecord record = importRecord(location, reference);
                imports.compute(record.key(), (ignored, existing) -> {
                    if (existing == null) {
                        return record;
                    }
                    existing.referenceTypes.add(referenceType(reference));
                    return existing;
                });
            }
        }
        List<ImportRecord> sorted = new ArrayList<>(imports.values());
        Collections.sort(sorted);
        List<Map<String, Object>> values = new ArrayList<>();
        for (ImportRecord record : sorted) {
            values.add(record.toJson());
        }
        return values;
    }

    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 ImportRecord importRecord(ExternalLocation location, Reference reference) {
        String name = firstNonEmpty(
            location.getOriginalImportedName(),
            location.getLabel(),
            location.getSymbol() == null ? null : location.getSymbol().getName()
        );
        if (name == null) {
            name = reference.getToAddress() == null ? "unknown" : reference.getToAddress().toString();
        }
        String library = emptyToNull(location.getLibraryName());
        Address address = location.getExternalSpaceAddress();
        if (address == null) {
            address = location.getAddress();
        }
        return new ImportRecord(name, library, importKind(location, reference), address, referenceType(reference));
    }

    private static String importKind(ExternalLocation location, Reference reference) {
        if (reference.getReferenceType().isCall() || location.isFunction()) {
            return "function";
        }
        if (reference.getReferenceType().isData()
            || reference.getReferenceType().isRead()
            || reference.getReferenceType().isWrite()) {
            return "data";
        }
        return "unknown";
    }

    private static String referenceType(Reference reference) {
        return reference.getReferenceType().getName();
    }

    private static List<Map<String, Object>> exports(Program program, Function function) {
        Address entry = function.getEntryPoint();
        if (!program.getSymbolTable().isExternalEntryPoint(entry)) {
            return Collections.emptyList();
        }

        Set<String> names = new TreeSet<>();
        SymbolIterator symbols = program.getSymbolTable().getSymbolsAsIterator(entry);
        while (symbols.hasNext()) {
            Symbol symbol = symbols.next();
            if (!symbol.isDynamic() && symbol.getName() != null && !symbol.getName().isBlank()) {
                names.add(symbol.getName());
            }
        }
        if (names.isEmpty()) {
            names.add(function.getName());
        }

        List<Map<String, Object>> values = new ArrayList<>();
        for (String name : names) {
            Map<String, Object> export = new LinkedHashMap<>();
            export.put("name", name);
            export.put("address", ProgramModel.address(entry));
            export.put("kind", "entry_point");
            values.add(export);
        }
        return values;
    }

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

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

    private static final class ImportRecord implements Comparable<ImportRecord> {
        private final String name;
        private final String library;
        private final String kind;
        private final Address address;
        private final Set<String> referenceTypes = new TreeSet<>();

        private ImportRecord(String name, String library, String kind, Address address, String referenceType) {
            this.name = name;
            this.library = library;
            this.kind = kind;
            this.address = address;
            this.referenceTypes.add(referenceType);
        }

        private String key() {
            return String.join("\u001f", name, nullToEmpty(library), kind, address == null ? "" : address.toString());
        }

        private Map<String, Object> toJson() {
            Map<String, Object> value = new LinkedHashMap<>();
            value.put("name", name);
            value.put("library", library);
            value.put("kind", kind);
            value.put("address", ProgramModel.address(address));
            value.put("reference_types", new ArrayList<>(referenceTypes));
            return value;
        }

        @Override
        public int compareTo(ImportRecord other) {
            return key().compareTo(other.key());
        }

        private static String nullToEmpty(String value) {
            return value == null ? "" : value;
        }
    }
}
