package org.ghidra.bridge.query;

import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import ghidra.program.model.address.Address;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.InstructionIterator;
import ghidra.program.model.listing.Program;
import org.ghidra.bridge.serialization.JsonWriter;

public final class ProgramSourceMaps {
    private ProgramSourceMaps() {
    }

    public static String sourceMapJson(Program program, String functionAddress) {
        Map<String, Object> identity = sourceMap(
            program,
            ProgramModel.requireFunctionAt(program, functionAddress)
        );
        return identity == null ? null : JsonWriter.json(identity);
    }

    public static Map<String, Object> sourceMap(Program program, Function function) {
        List<Object> entries = sourceMapEntries(program, function);
        if (entries.isEmpty()) {
            return null;
        }

        Map<String, List<Long>> grouped = new LinkedHashMap<>();
        String entrypointFile = sourceFilePath(entries.get(0));
        for (Object entry : entries) {
            String filePath = sourceFilePath(entry);
            if (filePath == null || filePath.isEmpty()) {
                continue;
            }
            grouped.putIfAbsent(filePath, new ArrayList<>());
            Long line = sourceLineNumber(entry);
            if (line != null) {
                grouped.get(filePath).add(line);
            }
        }
        if (grouped.isEmpty()) {
            return null;
        }

        List<String> candidates = new ArrayList<>(grouped.keySet());
        Collections.sort(candidates);
        String selected = null;
        for (String filePath : candidates) {
            if (selected == null || sourceFilePrecedes(filePath, selected, entrypointFile, grouped)) {
                selected = filePath;
            }
        }
        if (selected == null) {
            return null;
        }

        List<Long> lines = grouped.get(selected);
        Map<String, Object> identity = new LinkedHashMap<>();
        identity.put("file_path", selected);
        identity.put("function_name", function.getName());
        identity.put("line_start", lines.isEmpty() ? null : Collections.min(lines));
        identity.put("line_end", lines.isEmpty() ? null : Collections.max(lines));
        return identity;
    }

    private static boolean sourceFilePrecedes(
        String candidate,
        String selected,
        String entrypointFile,
        Map<String, List<Long>> grouped
    ) {
        boolean candidateIsEntry = candidate.equals(entrypointFile);
        boolean selectedIsEntry = selected.equals(entrypointFile);
        if (candidateIsEntry != selectedIsEntry) {
            return candidateIsEntry;
        }
        int candidateLineCount = grouped.get(candidate).size();
        int selectedLineCount = grouped.get(selected).size();
        if (candidateLineCount != selectedLineCount) {
            return candidateLineCount > selectedLineCount;
        }
        return candidate.compareTo(selected) < 0;
    }

    private static List<Object> sourceMapEntries(Program program, Function function) {
        List<Object> entries = new ArrayList<>();
        Object manager = sourceFileManager(program);
        if (manager == null) {
            return entries;
        }
        entries.addAll(sourceMapEntriesAt(manager, function.getEntryPoint()));
        InstructionIterator instructions = program.getListing().getInstructions(function.getBody(), true);
        while (instructions.hasNext()) {
            entries.addAll(sourceMapEntriesAt(manager, instructions.next().getAddress()));
        }
        return entries;
    }

    private static Object sourceFileManager(Program program) {
        try {
            Method method = program.getClass().getMethod("getSourceFileManager");
            return method.invoke(program);
        } catch (ReflectiveOperationException e) {
            return null;
        }
    }

    private static List<Object> sourceMapEntriesAt(Object manager, Address address) {
        List<Object> result = new ArrayList<>();
        try {
            Method method = manager.getClass().getMethod("getSourceMapEntries", Address.class);
            Object entries = method.invoke(manager, address);
            if (entries instanceof Iterable<?>) {
                for (Object entry : (Iterable<?>) entries) {
                    result.add(entry);
                }
            }
        } catch (ReflectiveOperationException e) {
            return result;
        }
        return result;
    }

    private static String sourceFilePath(Object entry) {
        try {
            Object sourceFile = entry.getClass().getMethod("getSourceFile").invoke(entry);
            if (sourceFile == null) {
                return null;
            }
            for (String methodName : new String[] {"getPath", "getFilename", "getName"}) {
                try {
                    Object value = sourceFile.getClass().getMethod(methodName).invoke(sourceFile);
                    if (value != null && !value.toString().isEmpty()) {
                        return value.toString();
                    }
                } catch (ReflectiveOperationException ignored) {
                }
            }
        } catch (ReflectiveOperationException e) {
            return null;
        }
        return null;
    }

    private static Long sourceLineNumber(Object entry) {
        for (String methodName : new String[] {"getLineNumber", "getLine"}) {
            try {
                Object value = entry.getClass().getMethod(methodName).invoke(entry);
                if (value instanceof Number) {
                    long line = ((Number) value).longValue();
                    return line > 0 ? line : null;
                }
            } catch (ReflectiveOperationException ignored) {
            }
        }
        return null;
    }
}
