import ghidra.app.script.GhidraScript;
import ghidra.app.decompiler.DecompInterface;
import ghidra.app.decompiler.DecompileResults;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.FunctionManager;
import ghidra.program.model.listing.Instruction;
import ghidra.program.model.address.Address;
import ghidra.program.model.address.AddressSpace;
import ghidra.program.model.address.AddressIterator;
import ghidra.program.model.symbol.Reference;
import ghidra.program.model.mem.Memory;
import java.io.PrintWriter;
import java.io.FileWriter;
import java.io.BufferedReader;
import java.io.FileReader;
import java.util.*;

/**
 * BFS-closure over the DSP call-graph starting from soothe2 DSP vtable slots and all
 * FUN_* in decomp_dsp.txt. Decompiles every reachable function (if missing) and dumps
 * all .data float constants referenced from instructions.
 */
public class DumpFuns extends GhidraScript {

    final long LO = 0x180000000L;
    final long HI = 0x182a00000L;

    long parseAddr(String s) {
        try {
            return Long.parseLong(s.replaceAll("[^0-9a-fA-F]", ""), 16);
        } catch (Exception e) { return -1; }
    }

    @Override
    public void run() throws Exception {
        FunctionManager fm = currentProgram.getFunctionManager();
        AddressSpace as = currentProgram.getAddressFactory().getDefaultAddressSpace();

        // --- seeds from decomp_dsp.txt ---
        Set<Long> seed = new LinkedHashSet<>();
        BufferedReader br = new BufferedReader(new FileReader("/home/m/re-tools/decomp_dsp.txt"));
        String line;
        while ((line = br.readLine()) != null) {
            int i = line.indexOf("FUN_");
            while (i >= 0) {
                int e = i + 4;
                while (e < line.length() && "0123456789abcdefABCDEF".indexOf(line.charAt(e)) >= 0) e++;
                long a = parseAddr(line.substring(i + 4, e));
                if (a >= LO && a <= HI) seed.add(a);
                i = line.indexOf("FUN_", e);
            }
        }
        br.close();
        // --- seeds: vtable slots ---
        long[] vts = {
            0x1824abb90L, 0x1824ac7a8L, 0x1824ab7c0L, 0x1824ac638L, 0x1824ac5d8L,
            0x1824be810L, 0x1824be228L, 0x1824be0a8L, 0x1824be7a0L, 0x1824ac210L,
            0x1824ac248L, 0x1824b1ac8L, 0x1824b11c8L, 0x1824b0fd0L, 0x1824b1178L
        };
        for (long base : vts) {
            for (int i = 0; i < 64; i++) {
                try {
                    Address va = as.getAddress(base + i * 8L);
                    long tgt = currentProgram.getMemory().getLong(va);
                    if (tgt >= LO && tgt <= HI) seed.add(tgt);
                } catch (Exception e) { }
            }
        }

        // --- BFS closure (no decompile needed for called-functions graph) ---
        Set<Long> done = new TreeSet<>();
        Deque<Long> queue = new ArrayDeque<>(seed);
        List<Long> order = new ArrayList<>();
        int cap = 5000;
        while (!queue.isEmpty() && order.size() < cap) {
            long a = queue.poll();
            if (!done.add(a)) continue;
            order.add(a);
            Function f = fm.getFunctionAt(as.getAddress(a));
            if (f == null) continue;
            try {
                for (Function c : f.getCalledFunctions(monitor)) {
                    long ca = c.getEntryPoint().getOffset();
                    if (ca >= LO && ca <= HI && !done.contains(ca)) queue.add(ca);
                }
            } catch (Exception e) { }
        }

        // --- fun_map always written fresh ---
        PrintWriter pm = new PrintWriter(new java.io.BufferedWriter(
            new FileWriter("/home/m/re-tools/fun_map.txt", false)));
        for (long a : order) {
            Function f = fm.getFunctionAt(as.getAddress(a));
            if (f == null) continue;
            pm.println(Long.toHexString(a) + " " + f.getBody().getNumAddresses() + " " + f.getName());
        }
        pm.close();

        // --- decompile missing ---
        java.io.File df = new java.io.File("/home/m/re-tools/decomp_funs.txt");
        boolean haveDec = df.exists() && df.length() > 100000;
        DecompInterface di = null;
        PrintWriter pw = null;
        if (!haveDec) {
            di = new DecompInterface();
            di.openProgram(currentProgram);
            pw = new PrintWriter(new java.io.BufferedWriter(new FileWriter(df)));
        }
        if (pw != null) {
            int n = 0;
            for (long a : order) {
                Function f = fm.getFunctionAt(as.getAddress(a));
                if (f == null) continue;
                if (f.getBody().getNumAddresses() < 30) continue;
                DecompileResults res = di.decompileFunction(f, 120, monitor);
                if (res != null && res.getDecompiledFunction() != null) {
                    n++;
                    pw.println("############ FUN_ " + Long.toHexString(a) + " size=" +
                        f.getBody().getNumAddresses() + " ############");
                    pw.println(res.getDecompiledFunction().getC());
                    pw.println();
                }
            }
            pw.close();
            di.dispose();
            println("DECOMP_DONE n=" + n);
        }

        // --- constants from instruction references ---
        Map<Long,String> consts = new TreeMap<>();
        Memory mem = currentProgram.getMemory();
        AddressIterator it = mem.getLoadedAndInitializedAddressSet().getAddresses(true);
        // iterate references: for data targets in range, collect functions referencing them
        ghidra.program.model.symbol.ReferenceManager rm = currentProgram.getReferenceManager();
        for (Function f : fm.getFunctions(true)) {
            long fa = f.getEntryPoint().getOffset();
            if (!order.contains(fa)) continue;
            Iterator<Instruction> iit = currentProgram.getListing().getInstructions(f.getBody(), true);
            while (iit.hasNext()) {
                Instruction ins = iit.next();
                for (Reference r : rm.getReferencesFrom(ins.getAddress())) {
                    long ta = r.getToAddress().getOffset();
                    if ((ta >= 0x1824c0000L && ta <= 0x182700000L) ||
                        (ta >= 0x180000000L && ta <= 0x181200000L)) {
                        consts.merge(ta, Long.toHexString(fa), (x, y) -> x + ";" + y);
                    }
                }
            }
        }
        PrintWriter pc = new PrintWriter(new java.io.BufferedWriter(
            new FileWriter("/home/m/re-tools/consts.txt", false)));
        for (Map.Entry<Long,String> e : consts.entrySet()) {
            try {
                long ta = e.getKey();
                int[] b = new int[8];
                boolean ok = true;
                for (int k = 0; k < 8; k++) {
                    Address a = as.getAddress(ta + k);
                    if (!mem.getLoadedAndInitializedAddressSet().contains(a)) { ok = false; break; }
                    b[k] = mem.getByte(a) & 0xff;
                }
                if (!ok) continue;
                float f32 = Float.intBitsToFloat(b[0] | (b[1] << 8) | (b[2] << 16) | (b[3] << 24));
                float f32b = Float.intBitsToFloat(b[4] | (b[5] << 8) | (b[6] << 16) | (b[7] << 24));
                pc.println(Long.toHexString(ta) + " f32=" + f32 + " f32b=" + f32b + "  refs=" + e.getValue());
            } catch (Exception ex) { }
        }
        pc.close();
        println("CONSTS_DONE nn=" + order.size() + " consts=" + consts.size());
    }
}