package org.ghidra.bridge.query;

import java.util.ArrayList;
import java.util.IdentityHashMap;
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.Data;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.FunctionIterator;
import ghidra.program.model.listing.Parameter;
import ghidra.program.model.listing.Program;
import ghidra.program.model.symbol.Symbol;
import ghidra.program.model.symbol.SymbolIterator;
import org.ghidra.bridge.serialization.JsonWriter;
import org.ghidra.bridge.serialization.ProgramTypes;

public final class ProgramMetadata {
    private ProgramMetadata() {
    }

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

    public static Map<String, Object> metadata(Program program) {
        return metadata(program, new ProgramTypes());
    }

    public static Map<String, Object> metadata(Program program, ProgramTypes types) {
        SymbolTable symbols = collectSymbols(program, types);
        List<Map<String, Object>> functions = collectFunctions(program, types, symbols);

        Map<String, Object> value = new LinkedHashMap<>();
        value.put("symbols", symbols.records());
        value.put("functions", functions);
        value.put("types", types.records());
        return value;
    }

    private static SymbolTable collectSymbols(Program program, ProgramTypes types) {
        List<SymbolRecord> records = new ArrayList<>();
        SymbolIterator iterator = program.getSymbolTable().getSymbolIterator();
        while (iterator.hasNext()) {
            Symbol symbol = iterator.next();
            if (symbol.getAddress() == null) {
                continue;
            }
            records.add(symbolRecord(program, types, symbol));
        }
        records.sort((left, right) -> left.key().compareTo(right.key()));

        Map<String, Integer> counts = new LinkedHashMap<>();
        Map<Symbol, String> ids = new IdentityHashMap<>();
        Map<String, String> idsByFunction = new LinkedHashMap<>();
        List<Map<String, Object>> values = new ArrayList<>();
        for (SymbolRecord record : records) {
            String base = "symbol:" + ProgramTypes.idText(record.key());
            int count = counts.getOrDefault(base, 0);
            counts.put(base, count + 1);
            String id = count == 0 ? base : base + "#" + count;
            ids.put(record.symbol, id);
            if (record.functionEntry != null && record.primary) {
                idsByFunction.put(record.functionEntry.toString(), id);
            }
            values.add(record.toJson(id));
        }
        return new SymbolTable(values, ids, idsByFunction);
    }

    private static SymbolRecord symbolRecord(Program program, ProgramTypes types, Symbol symbol) {
        Address address = symbol.getAddress();
        Function function = program.getFunctionManager().getFunctionAt(address);
        String typeId = null;
        if (function == null) {
            Data data = program.getListing().getDataAt(address);
            if (data == null) {
                data = program.getListing().getDataContaining(address);
            }
            if (data != null) {
                typeId = types.id(data.getDataType());
            }
        }
        return new SymbolRecord(
            symbol,
            symbol.getName(),
            symbolKind(symbol.getSymbolType().toString()),
            symbolSource(symbol.getSource().toString()),
            address,
            namespace(symbol),
            symbol.isPrimary(),
            symbol.isExternal(),
            typeId,
            function == null ? null : function.getEntryPoint()
        );
    }

    private static List<Map<String, Object>> collectFunctions(
        Program program,
        ProgramTypes types,
        SymbolTable symbols
    ) {
        List<Map<String, Object>> values = new ArrayList<>();
        FunctionIterator iterator = program.getFunctionManager().getFunctions(true);
        while (iterator.hasNext()) {
            Function function = iterator.next();
            if (!function.isExternal()) {
                values.add(functionRecord(types, symbols, function));
            }
        }
        values.sort((left, right) -> stringValue(left, "id").compareTo(stringValue(right, "id")));
        return values;
    }

    private static Map<String, Object> functionRecord(
        ProgramTypes types,
        SymbolTable symbols,
        Function function
    ) {
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("id", functionId(function.getEntryPoint()));
        value.put("name", function.getName());
        value.put("entry", ProgramModel.address(function.getEntryPoint()));
        value.put("namespace", namespace(function));
        value.put("external", function.isExternal());
        value.put("thunk", function.isThunk());
        value.put("thunked_function_id", thunkedFunctionId(function));
        value.put("symbol_id", symbols.idForFunction(function));
        value.put("signature", signature(types, function));
        return value;
    }

    private static Map<String, Object> signature(ProgramTypes types, Function function) {
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("display", function.getSignature().toString());
        value.put("calling_convention", emptyToFallback(function.getCallingConventionName(), "unknown"));
        value.put("return_type_id", types.id(function.getReturnType()));
        List<Map<String, Object>> parameters = new ArrayList<>();
        for (Parameter parameter : function.getParameters()) {
            Map<String, Object> record = new LinkedHashMap<>();
            record.put("ordinal", parameter.getOrdinal());
            record.put("name", emptyToFallback(parameter.getName(), "param_" + parameter.getOrdinal()));
            record.put("type_id", types.id(parameter.getDataType()));
            record.put(
                "storage",
                parameter.getVariableStorage() == null ? "unknown" : parameter.getVariableStorage().toString()
            );
            parameters.add(record);
        }
        value.put("parameters", parameters);
        value.put("varargs", function.hasVarArgs());
        value.put("no_return", function.hasNoReturn());
        return value;
    }

    private static String thunkedFunctionId(Function function) {
        if (!function.isThunk()) {
            return null;
        }
        Function thunked = function.getThunkedFunction(true);
        return thunked == null || thunked.isExternal() ? null : functionId(thunked.getEntryPoint());
    }

    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";
        }
    }

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

    private static String namespace(Symbol symbol) {
        return symbol.getParentNamespace() == null ? "" : symbol.getParentNamespace().toString();
    }

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

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

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

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

    private static final class SymbolTable {
        private final List<Map<String, Object>> records;
        private final Map<Symbol, String> ids;
        private final Map<String, String> idsByFunction;

        private SymbolTable(
            List<Map<String, Object>> records,
            Map<Symbol, String> ids,
            Map<String, String> idsByFunction
        ) {
            this.records = records;
            this.ids = ids;
            this.idsByFunction = idsByFunction;
        }

        private List<Map<String, Object>> records() {
            return records;
        }

        private String idForFunction(Function function) {
            String id = ids.get(function.getSymbol());
            return id != null ? id : idsByFunction.get(function.getEntryPoint().toString());
        }
    }

    private static final class SymbolRecord {
        private final Symbol symbol;
        private final String name;
        private final String kind;
        private final String source;
        private final Address address;
        private final String namespace;
        private final boolean primary;
        private final boolean external;
        private final String typeId;
        private final Address functionEntry;

        private SymbolRecord(
            Symbol symbol,
            String name,
            String kind,
            String source,
            Address address,
            String namespace,
            boolean primary,
            boolean external,
            String typeId,
            Address functionEntry
        ) {
            this.symbol = symbol;
            this.name = name;
            this.kind = kind;
            this.source = source;
            this.address = address;
            this.namespace = namespace;
            this.primary = primary;
            this.external = external;
            this.typeId = typeId;
            this.functionEntry = functionEntry;
        }

        private String key() {
            return String.join(
                "\u001f",
                namespace,
                name,
                kind,
                address.toString(),
                source,
                Boolean.toString(primary),
                Boolean.toString(external)
            );
        }

        private Map<String, Object> toJson(String id) {
            Map<String, Object> value = new LinkedHashMap<>();
            value.put("id", id);
            value.put("name", name);
            value.put("kind", kind);
            value.put("source", source);
            value.put("address", ProgramModel.address(address));
            value.put("namespace", namespace);
            value.put("primary", primary);
            value.put("external", external);
            value.put("type_id", typeId);
            return value;
        }
    }
}
