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.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.*;

/**
 * DumpFuns2 - unfrozen BFS decompiler.
 * Differences vs DumpFuns.java:
 *   1. NO haveDec gate: always decompiles into decomp_funs2.txt.
 *   2. Extended seeds: explicit critical DSP addresses + reverse callers
 *      (references-to) so indirect/vtable-dispatched functions enter closure.
 *   3. Skip threshold lowered from 30 to 8 instructions (keeps dispatchers/stubs).
 * Also writes fun_map2.txt (fresh) and re-scans .data float constants.
 */
public class DumpFuns2 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; }
    }

    boolean inRange(long a) { return a >= LO && a <= HI; }

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

        // --- seeds: all FUN_ from decomp_dsp.txt + explicit critical DSP addrs ---
        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 (inRange(a)) seed.add(a);
                i = line.indexOf("FUN_", e);
            }
        }
        br.close();

        long[] critical = {
            0x180529fe0L, // mask/coefficient setup (vtable-only, 2051 instr)
            0x180563440L, // LUT curve + gamma + combine
            0x180563ce0L, // IIR level-tracker INIT
            0x18056e3e0L, // twin-mask factory
            0x18052f500L, // interleave
            0x18052b550L, // FFT-conv loop
            0x18052ee70L, // per-bin gain
            0x18052d650L, // band setup
            0x18052d920L, // window
            0x18052e190L, // buffer alloc
            0x18052dc30L, // FFT plan builder
            0x18052da00L, // ramp fill
            0x18052dbc0L, // cplx-interleave
            0x180535880L, // twin kernel
            0x180536300L, // twin caller
            0x180536f90L, // dead sibling
            0x180530d30L, // level-weight formula
            0x1805316e0L, // coeff writer (17 case)
            0x180563fa0L, // IIR level-tracker UPDATE (caller chain 5631c0)
            0x1805631c0L, // LUT+IIR-init+UPDATE dispatcher
            0x180563260L, // param dispatcher
            0x180563a60L  // combine (called from 563440)
        };
        for (long a : critical) seed.add(a);

        // --- 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 (inRange(tgt)) seed.add(tgt);
                } catch (Exception e) { }
            }
        }

        // --- BFS closure forward + REVERSE (indirect callers) ---
        Set<Long> done = new TreeSet<>();
        Deque<Long> queue = new ArrayDeque<>(seed);
        List<Long> order = new ArrayList<>();
        int cap = 8000;
        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 (inRange(ca) && !done.contains(ca)) queue.add(ca);
                }
            } catch (Exception e) { }
            // reverse: any function referencing this entry (indirect dispatch callers)
            try {
                for (Reference r : currentProgram.getReferenceManager()
                        .getReferencesTo(as.getAddress(a))) {
                    Address from = r.getFromAddress();
                    if (from == null) continue;
                    Function cf = fm.getFunctionContaining(from);
                    if (cf != null) {
                        long ca = cf.getEntryPoint().getOffset();
                        if (inRange(ca) && !done.contains(ca)) queue.add(ca);
                    }
                }
            } catch (Exception e) { }
        }

        // --- fun_map2.txt fresh ---
        PrintWriter pm = new PrintWriter(new java.io.BufferedWriter(
            new FileWriter("/home/m/re-tools/fun_map2.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 ALWAYS into decomp_funs2.txt ---
        DecompInterface di = new DecompInterface();
        di.openProgram(currentProgram);
        PrintWriter pw = new PrintWriter(new java.io.BufferedWriter(
            new FileWriter("/home/m/re-tools/decomp_funs2.txt", false)));
        int n = 0, skipped = 0;
        for (long a : order) {
            Function f = fm.getFunctionAt(as.getAddress(a));
            if (f == null) continue;
            long sz = f.getBody().getNumAddresses();
            if (sz < 8) { skipped++; continue; }
            DecompileResults res = di.decompileFunction(f, 120, monitor);
            if (res != null && res.getDecompiledFunction() != null) {
                n++;
                pw.println("############ FUN_ " + Long.toHexString(a) + " size=" + sz + " ############");
                pw.println(res.getDecompiledFunction().getC());
                pw.println();
            }
        }
        pw.close();
        di.dispose();
        println("DECOMP2_DONE n=" + n + " skipped=" + skipped + " order=" + order.size());

        // --- constants (same as DumpFuns, scaled to what we touched) ---
        Map<Long,String> consts = new TreeMap<>();
        Memory mem = currentProgram.getMemory();
        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 : currentProgram.getReferenceManager().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/consts2.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("CONSTS2_DONE consts=" + consts.size());
    }
}