import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.*;
import ghidra.program.model.mem.*;
import ghidra.program.model.listing.*;
import ghidra.program.model.symbol.*;
import ghidra.app.decompiler.*;
import java.util.*;
import java.util.regex.*;
import java.io.*;

// Pass 1 (read-only): anchors for the turn/main-loop spine.
//  A. strings containing SE*/SNM*/AppStartup/etc. -> xrefs -> containing functions (+ their callers)
//  B. symbols (RTTI/vftables) containing those names
//  C. entry point + callees (WinMain hunt), CreateThread sites + thread procs
//  D. std::string layout check (FUN_008b9d70)
public class SpineRecon extends GhidraScript {
    DecompInterface decomp; Memory mem; PrintWriter out; ReferenceManager rm; SymbolTable st;
    static final String[] NEEDLES = {
        "SETurnEndPending","SNMEndTurn","SEProcessTurn","SNMDoEncounterQuery","SNMAllCombatDone","SNMResumePlaying",
        "SEAIPrepareTurn","SNMRunAI","SEFleetArrived","SEBuildCompleted","SETurnResults","SETurnEvents","AppStartup",
        "SETurnTimeExpired","SEResumePlaying","SNMEncounterQueryResults","SNMHostCombat","SNMLaunchCombat","SNMEncounterResults",
        "SETurnEndCancelled","SNMQueryEndTurnDone","SELastPlaying","PROCESSING_TITLE","Initializing Game","LOBBYSTATUS_STRATEGY_ROUND",
        "SENewTurn","SEStartTurn","SEEndTurn","StrategyServer","TurnCommands","ProcessTurn","EndTurn","MainLoop","Tick","AppMain"
    };
    static final String[] DECOMP = {"00902350","00736e30","008a0e50","008b9d70","008b9d90"};

    String cstr(Address a, int max) {
        try {
            byte[] b = new byte[max]; int got = mem.getBytes(a, b);
            int i = 0; while (i < got && b[i] != 0 && (b[i]&0xff) >= 0x20 && (b[i]&0xff) < 0x7f) i++;
            if (i >= 1 && i < got && b[i] == 0) return new String(b, 0, i, "ISO-8859-1");
        } catch (Exception e) {}
        return null;
    }
    String decompRes(Function f) {
        try {
            DecompileResults res = decomp.decompileFunction(f, 180, monitor);
            if (res == null || !res.decompileCompleted()) return "[decompile failed]";
            String c = res.getDecompiledFunction().getC();
            Matcher m = Pattern.compile("&?(DAT|PTR_s_|s_[A-Za-z0-9_]*|PTR_DAT|u_[A-Za-z0-9_]*)_([0-9a-f]{8})").matcher(c);
            StringBuffer sb = new StringBuffer();
            while (m.find()) {
                Address a = toAddr(Long.parseLong(m.group(2), 16));
                String s = cstr(a, 64); String rep = m.group(0);
                if (s != null && s.length() <= 60) rep = "\"" + s + "\"";
                m.appendReplacement(sb, Matcher.quoteReplacement(rep));
            }
            m.appendTail(sb); return sb.toString();
        } catch (Exception e) { return "[exception " + e.getMessage() + "]"; }
    }
    void dumpFunc(Function f, String tag) throws IOException {
        String n = String.format("%08x", f.getEntryPoint().getOffset());
        File fl = new File("/tmp/spine/" + n + ".c"); if (fl.exists()) return;
        PrintWriter w = new PrintWriter(new FileWriter(fl));
        w.println("// " + f.getName() + " @ " + f.getEntryPoint() + " size=" + f.getBody().getNumAddresses() + " tag=" + tag);
        StringBuilder cs = new StringBuilder();
        for (Function c : f.getCallingFunctions(monitor)) cs.append(c.getName() + "@" + c.getEntryPoint() + " ");
        w.println("// CALLERS: " + cs);
        StringBuilder ce = new StringBuilder();
        for (Function c : f.getCalledFunctions(monitor)) ce.append(c.getName() + " ");
        w.println("// CALLEES: " + ce);
        w.println(decompRes(f)); w.close();
    }
    String funcDesc(Function f) {
        StringBuilder cs = new StringBuilder();
        int n = 0;
        for (Function c : f.getCallingFunctions(monitor)) { if (n++ < 8) cs.append(c.getName() + " "); }
        return f.getName() + "@" + f.getEntryPoint() + " sz=" + f.getBody().getNumAddresses() + " callers(" + n + "):[" + cs.toString().trim() + "]";
    }

    @Override
    public void run() throws Exception {
        mem = currentProgram.getMemory(); rm = currentProgram.getReferenceManager(); st = currentProgram.getSymbolTable();
        decomp = new DecompInterface(); decomp.openProgram(currentProgram);
        new File("/tmp/spine").mkdirs();
        out = new PrintWriter(new FileWriter("/tmp/spine/recon.txt"));
        Set<Function> toDump = new LinkedHashSet<Function>();

        // A. strings
        out.println("==== A. STRING ANCHORS ====");
        for (MemoryBlock blk : mem.getBlocks()) {
            if (!blk.isInitialized() || blk.isExecute()) continue;
            if (!blk.getName().equals(".rdata") && !blk.getName().equals(".data")) continue;
            int len = (int) blk.getSize(); byte[] buf = new byte[len]; blk.getBytes(blk.getStart(), buf);
            int start = 0;
            for (int i = 0; i < len; i++) {
                if (buf[i] != 0) continue;
                int slen = i - start;
                if (slen >= 4 && slen <= 120) {
                    boolean ascii = true;
                    for (int k = start; k < i; k++) { int c = buf[k] & 0xff; if (c < 0x20 || c > 0x7e) { ascii = false; break; } }
                    if (ascii) {
                        String s = new String(buf, start, slen, "ISO-8859-1");
                        for (String nd : NEEDLES) {
                            if (!s.contains(nd)) continue;
                            Address a = blk.getStart().add(start);
                            StringBuilder sb = new StringBuilder();
                            ReferenceIterator ri = rm.getReferencesTo(a); int n = 0;
                            while (ri.hasNext()) {
                                Reference r = ri.next(); n++;
                                Function f = getFunctionContaining(r.getFromAddress());
                                if (f == null) { sb.append("\n      ?" + r.getFromAddress()); continue; }
                                sb.append("\n      " + r.getFromAddress() + " in " + funcDesc(f));
                                if (nd.startsWith("SE") || nd.startsWith("SNM") || nd.equals("AppStartup") || nd.equals("PROCESSING_TITLE") || nd.equals("Initializing Game")) toDump.add(f);
                            }
                            out.println(a + "  \"" + s + "\"  [" + nd + "] refs=" + n + sb);
                            break;
                        }
                    }
                }
                start = i + 1;
            }
        }
        // B. symbols
        out.println("\n==== B. SYMBOLS ====");
        SymbolIterator si = st.getAllSymbols(true);
        while (si.hasNext()) {
            Symbol s = si.next(); String n = s.getName(true);
            boolean hit = false;
            for (String nd : NEEDLES) if (n.contains(nd)) { hit = true; break; }
            if (!hit) continue;
            if (n.contains("StreamableHelper") || n.contains("VectorHelper")) continue;
            out.println(s.getAddress() + "  " + n + "  [" + s.getSymbolType() + "]");
            if (n.endsWith("::vftable")) {
                Address a = s.getAddress();
                for (int i = 0; i < 12; i++) {
                    try {
                        long p = mem.getInt(a.add(i*4)) & 0xffffffffL; Address fa = toAddr(p);
                        Function f = getFunctionAt(fa);
                        if (f == null) { if (i > 0) break; out.println("      [" + i + "] " + fa + " (nofunc)"); continue; }
                        out.println("      [" + i + "] " + funcDesc(f));
                        if (f.getBody().getNumAddresses() > 40) toDump.add(f);
                    } catch (Exception e) { break; }
                }
            }
        }
        // C. entry + CreateThread
        out.println("\n==== C. ENTRY / THREADS ====");
        for (Symbol s : st.getSymbols("entry")) {
            Function f = getFunctionAt(s.getAddress());
            out.println("entry sym @ " + s.getAddress() + " func=" + (f == null ? "null" : funcDesc(f)));
            if (f != null) {
                toDump.add(f);
                for (Function c : f.getCalledFunctions(monitor)) {
                    out.println("   entry callee: " + funcDesc(c) + " params=" + c.getParameterCount());
                    if (c.getBody().getNumAddresses() > 200 && !c.getName().startsWith("_")) toDump.add(c);
                }
            }
        }
        for (Symbol s : st.getSymbols("CreateThread")) {
            out.println("CreateThread sym @ " + s.getAddress() + " type=" + s.getSymbolType());
            ReferenceIterator ri = rm.getReferencesTo(s.getAddress());
            while (ri.hasNext()) {
                Reference r = ri.next(); Function f = getFunctionContaining(r.getFromAddress());
                out.println("   ref " + r.getFromAddress() + " " + r.getReferenceType() + " in " + (f == null ? "?" : funcDesc(f)));
                if (f != null) toDump.add(f);
                // thunk? follow refs to the thunk
                if (f != null && f.isThunk() || (f != null && f.getBody().getNumAddresses() < 12)) {
                    ReferenceIterator ri2 = rm.getReferencesTo(f.getEntryPoint());
                    while (ri2.hasNext()) { Reference r2 = ri2.next(); Function f2 = getFunctionContaining(r2.getFromAddress());
                        out.println("      thunk-ref " + r2.getFromAddress() + " in " + (f2 == null ? "?" : funcDesc(f2))); if (f2 != null) toDump.add(f2); }
                }
            }
        }
        // thread procs: look at the CreateThread call sites listed in prior recon; find the pushed function pointer
        long[] sites = {0x0090242dL, 0x008a14efL, 0x00736e84L};
        for (long sa : sites) {
            Address a = toAddr(sa);
            out.println("site " + a + ":");
            Instruction ins = getInstructionAt(a);
            int back = 0;
            Instruction p = ins;
            while (p != null && back < 12) {
                p = p.getPrevious(); back++;
                if (p == null) break;
                StringBuilder refs = new StringBuilder();
                for (Reference r : p.getReferencesFrom()) {
                    Function tf = getFunctionAt(r.getToAddress());
                    if (tf != null) { refs.append(" -> FUNC " + funcDesc(tf)); toDump.add(tf); }
                }
                out.println("   " + p.getAddress() + "  " + p + refs);
            }
        }
        for (String d : DECOMP) { Function f = getFunctionAt(toAddr(Long.parseLong(d, 16))); if (f != null) toDump.add(f); }
        // D. dump
        out.println("\n==== D. DUMPED ====");
        for (Function f : toDump) { out.println(f.getName() + " @ " + f.getEntryPoint()); dumpFunc(f, "spine"); }
        out.close(); decomp.dispose(); println("done");
    }
}