// ExportGaps -- the #60 re-export: FILL the dark regions, then decompile everything.
//
// Same output contract as ExportAll.java (so tools/gapcensus.py, the KB's @ADDR
// citations and every existing grep keep working):
//     <out>/all/part_NNN.c     "/* @ADDR file=TAG name=NAME */" + pseudocode, 400 fns/shard
//     <out>/functions_index.tsv   addr<TAB>size<TAB>fileTag<TAB>name
//     <out>/vtables.tsv           runs of >=3 consecutive function pointers in data
//     <out>/gapfill_report.tsv    what this pass CREATED (addr, evidence, round)
//
// What it adds: before exporting, force disassembly + function creation at every
// plausible function start Ghidra's auto-analysis missed --
//     (a) E8 call targets that land outside any function,
//     (b) data-section dwords pointing into .text (vtables / handler tables /
//         Performance pointers) at a plausible prologue,
//     (c) the starts discovered by tools/gapcensus.py (gap_census.tsv, col 5),
// iterated to a fixpoint because each new function reveals new call targets.
//
// Output dir is overridable:  analyzeHeadless ... -postScript ExportGaps.java <outDir>
import ghidra.app.script.GhidraScript;
import ghidra.app.decompiler.*;
import ghidra.program.model.address.*;
import ghidra.program.model.listing.*;
import ghidra.program.model.mem.*;
import ghidra.program.model.symbol.*;
import java.io.*;
import java.util.*;

public class ExportGaps extends GhidraScript {

    static final String DEFAULT_OUT = "C:\\git\\bt411\\reference\\decomp\\reexport";
    static final String CENSUS_TSV  = "C:\\git\\bt411\\reference\\decomp\\gap_census.tsv";

    FunctionManager fm;
    Memory mem;

    public void run() throws Exception {
        String[] args = getScriptArgs();
        File outDir = new File(args.length > 0 ? args[0] : DEFAULT_OUT);
        outDir.mkdirs();
        File allDir = new File(outDir, "all");
        allDir.mkdirs();

        fm = currentProgram.getFunctionManager();
        mem = currentProgram.getMemory();
        ReferenceManager rm = currentProgram.getReferenceManager();
        Listing listing = currentProgram.getListing();

        int before = fm.getFunctionCount();
        println("functions after auto-analysis: " + before);

        // ---------------------------------------------------------------
        // 0) THE GAP FILL
        // ---------------------------------------------------------------
        PrintWriter gf = new PrintWriter(new FileWriter(new File(outDir, "gapfill_report.tsv")));
        gf.println("addr\tevidence\tround");

        List<MemoryBlock> code = new ArrayList<>();
        List<MemoryBlock> dataBlocks = new ArrayList<>();
        for (MemoryBlock b : mem.getBlocks()) {
            if (!b.isInitialized()) continue;
            if (b.isExecute()) code.add(b); else dataBlocks.add(b);
        }

        // candidate -> evidence string (first evidence wins)
        LinkedHashMap<Long, String> cand = new LinkedHashMap<>();

        // (c) census starts -- cheap, load first so they carry that evidence tag
        int censusN = 0;
        File cf = new File(CENSUS_TSV);
        if (cf.exists()) {
            BufferedReader br = new BufferedReader(new FileReader(cf));
            String line = br.readLine();                      // header
            while ((line = br.readLine()) != null) {
                String[] f = line.split("\t", -1);
                if (f.length < 5 || f[4].isEmpty()) continue;
                for (String s : f[4].split(",")) {
                    if (s.isEmpty()) continue;
                    try { cand.put(Long.parseLong(s.trim(), 16), "census"); censusN++; }
                    catch (NumberFormatException ex) { /* skip */ }
                }
            }
            br.close();
        }
        println("census starts loaded: " + censusN);

        // (b) data-section pointers into code
        int dptrN = 0;
        for (MemoryBlock b : dataBlocks) {
            byte[] buf = new byte[(int) b.getSize()];
            try { b.getBytes(b.getStart(), buf); } catch (Exception ex) { continue; }
            for (int i = 0; i + 4 <= buf.length; i += 4) {
                long v = le32(buf, i);
                if (!inCode(code, v)) continue;
                if (!cand.containsKey(v)) { cand.put(v, "dataptr"); dptrN++; }
            }
        }
        println("data->code pointers: " + dptrN);

        // (a) E8 call targets (rounds happen below; the scan itself is per-round
        //     because newly disassembled bytes can expose more calls)
        int created = 0;
        for (int round = 1; round <= 4 && !monitor.isCancelled(); round++) {
            int callN = 0;
            for (MemoryBlock b : code) {
                byte[] buf = new byte[(int) b.getSize()];
                try { b.getBytes(b.getStart(), buf); } catch (Exception ex) { continue; }
                long base = b.getStart().getOffset();
                for (int i = 0; i + 5 <= buf.length; i++) {
                    if ((buf[i] & 0xff) != 0xE8) continue;
                    long tgt = base + i + 5 + (int) le32(buf, i + 1);
                    if (!inCode(code, tgt)) continue;
                    if (!cand.containsKey(tgt)) { cand.put(tgt, "call"); callN++; }
                }
            }
            println("round " + round + ": call targets added " + callN
                    + " (candidates " + cand.size() + ")");

            int made = 0;
            for (Map.Entry<Long, String> e : new ArrayList<>(cand.entrySet())) {
                if (monitor.isCancelled()) break;
                long va = e.getKey();
                Address a;
                try { a = toAddr(va); } catch (Exception ex) { continue; }
                if (fm.getFunctionContaining(a) != null) continue;      // already covered
                if (!"call".equals(e.getValue()) && !looksLikePrologue(a)) continue;
                try {
                    if (listing.getInstructionAt(a) == null) disassemble(a);
                    Function nf = createFunction(a, null);
                    if (nf != null) {
                        made++; created++;
                        gf.println(a + "\t" + e.getValue() + "\t" + round);
                    }
                } catch (Exception ex) { /* not code / collides -- skip */ }
            }
            println("round " + round + ": functions created " + made);
            if (made == 0) break;
        }
        gf.close();
        println("GAP FILL created " + created + " functions ("
                + before + " -> " + fm.getFunctionCount() + ")");

        // let the analyzers process the new code before decompiling
        try { analyzeChanges(currentProgram); } catch (Exception ex) {
            println("analyzeChanges skipped: " + ex.getMessage());
        }

        // ---------------------------------------------------------------
        // 1) file attribution via embedded assert paths (unchanged from ExportAll)
        // ---------------------------------------------------------------
        Map<Long, String> funcFile = new HashMap<>();
        DataIterator di = listing.getDefinedData(true);
        while (di.hasNext() && !monitor.isCancelled()) {
            Data d = di.next();
            if (d == null || !d.hasStringValue()) continue;
            Object v = d.getValue();
            if (v == null) continue;
            String s = v.toString().toLowerCase().replace('/', '\\');
            if (!s.contains("\\tesla_bt\\")) continue;
            if (!(s.endsWith(".cpp") || s.endsWith(".hpp"))) continue;
            String file = s.substring(s.lastIndexOf('\\') + 1);
            String dir = s.contains("\\bt_l4\\") ? "bt_l4" : s.contains("\\bt\\") ? "bt"
                       : s.contains("\\munga_l4\\") ? "munga_l4" : "munga";
            String tag = dir + "/" + file;
            for (Reference ref : rm.getReferencesTo(d.getAddress())) {
                Function f = fm.getFunctionContaining(ref.getFromAddress());
                if (f != null) funcFile.put(f.getEntryPoint().getOffset(), tag);
            }
        }
        println("attributed functions: " + funcFile.size());

        // ---------------------------------------------------------------
        // 2) vtable recovery (unchanged)
        // ---------------------------------------------------------------
        PrintWriter vt = new PrintWriter(new FileWriter(new File(outDir, "vtables.tsv")));
        int vtCount = 0;
        for (MemoryBlock b : mem.getBlocks()) {
            if (!b.isInitialized() || b.isExecute()) continue;
            Address a = b.getStart(), end = b.getEnd();
            List<Long> run = new ArrayList<>();
            Address runStart = null;
            while (a.compareTo(end) < 0 && !monitor.isCancelled()) {
                long val;
                try { val = mem.getInt(a) & 0xffffffffL; } catch (Exception ex) { break; }
                Function f = null;
                try { f = fm.getFunctionAt(a.getNewAddress(val)); } catch (Exception ex) {}
                if (f != null) { if (run.isEmpty()) runStart = a; run.add(val); }
                else { if (run.size() >= 3) { writeVt(vt, runStart, run); vtCount++; } run.clear(); }
                a = a.add(4);
            }
            if (run.size() >= 3) { writeVt(vt, runStart, run); vtCount++; }
        }
        vt.close();
        println("vtables recovered: " + vtCount);

        // ---------------------------------------------------------------
        // 3) decompile everything, sharded 400/file (unchanged contract)
        // ---------------------------------------------------------------
        DecompInterface dec = new DecompInterface();
        dec.openProgram(currentProgram);
        PrintWriter idx = new PrintWriter(new FileWriter(new File(outDir, "functions_index.tsv")));
        FunctionIterator it = fm.getFunctions(true);
        int count = 0, shard = 0, failed = 0;
        PrintWriter out = null;
        while (it.hasNext() && !monitor.isCancelled()) {
            Function f = it.next();
            if (count % 400 == 0) {
                if (out != null) out.close();
                out = new PrintWriter(new FileWriter(new File(allDir, String.format("part_%03d.c", shard++))));
            }
            count++;
            long off = f.getEntryPoint().getOffset();
            String tag = funcFile.getOrDefault(off, "?");
            idx.println(Long.toHexString(off) + "\t" + f.getBody().getNumAddresses() + "\t" + tag + "\t" + f.getName());
            try {
                DecompileResults r = dec.decompileFunction(f, 45, monitor);
                out.println("/* @" + f.getEntryPoint() + " file=" + tag + " name=" + f.getName() + " */");
                if (r != null && r.decompileCompleted() && r.getDecompiledFunction() != null)
                    out.println(r.getDecompiledFunction().getC());
                else { out.println("// <decompile failed>"); failed++; }
                out.println();
            } catch (Exception ex) {
                out.println("// <exception: " + ex.getMessage() + ">\n");
                failed++;
            }
            if (count % 500 == 0) println("decompiled " + count + " ...");
        }
        if (out != null) out.close();
        idx.close();
        dec.dispose();
        println("DONE. functions: " + count + " (shards " + shard + ", decompile failures " + failed + ")");
    }

    boolean inCode(List<MemoryBlock> code, long va) {
        for (MemoryBlock b : code)
            if (va >= b.getStart().getOffset() && va < b.getEnd().getOffset()) return true;
        return false;
    }

    boolean looksLikePrologue(Address a) {
        try {
            byte[] p = new byte[3];
            mem.getBytes(a, p);
            int b0 = p[0] & 0xff, b1 = p[1] & 0xff;
            return (b0 == 0x55 && b1 == 0x8b)        // push ebp; mov ebp,esp
                || b0 == 0x53 || b0 == 0x56 || b0 == 0x57   // push ebx/esi/edi
                || b0 == 0xC8                         // enter
                || (b0 == 0x83 && b1 == 0xEC)         // sub esp, imm8
                || (b0 == 0x81 && b1 == 0xEC);        // sub esp, imm32
        } catch (Exception ex) { return false; }
    }

    static long le32(byte[] b, int i) {
        return ((long) (b[i] & 0xff))
             | ((long) (b[i + 1] & 0xff) << 8)
             | ((long) (b[i + 2] & 0xff) << 16)
             | ((long) (b[i + 3] & 0xff) << 24);
    }

    void writeVt(PrintWriter vt, Address start, List<Long> run) {
        StringBuilder sb = new StringBuilder(start.toString());
        for (Long v : run) sb.append("\t").append(Long.toHexString(v));
        vt.println(sb.toString());
    }
}