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.program.model.data.*;
import ghidra.app.decompiler.*;
import java.util.*;
import java.util.regex.*;
import java.io.*;

// Fix-up: make the class structs Ghidra uses for __thiscall 'this' (VariableUtilities.findOrCreateClassStruct) BE the recovered
// structs; give the COL!=0 serializers custom storage with the shifted view; re-run verification decompiles.
public class FixThis extends GhidraScript {
    DataTypeManager dtm; SymbolTable st; PrintWriter log; DecompInterface decomp; Memory mem;
    static final String[][] CLASSES = { // {namespace path, /SOTS type name}
        {"Game::ServerSystem","Game::ServerSystem"},{"Game::ServerPlayer","Game::ServerPlayer"},{"Game::StarFleet","Game::StarFleet"},{"Game::StarShip","Game::StarShip"},{"Game::StrategyServer","Game::StrategyServer"},
        {"Game::Population","Game::Population"},{"Game::PopulationGroup","Game::PopulationGroup"},{"Game::Morale","Game::Morale"},{"Game::MoraleEvent","Game::MoraleEvent"},{"Game::IndependenceInfo","Game::IndependenceInfo"},
        {"Game::StarSystem::OutputRates","Game::StarSystem::OutputRates"},{"Game::StarSystem::PlayerView","Game::StarSystem::PlayerView"},{"Game::ShipBuildOrder","Game::ShipBuildOrder"},{"Game::DiplomacyStats","Game::DiplomacyStats"},
        {"Game::PlayerReport","Game::PlayerReport"},{"Game::PlayerAlliances","Game::PlayerAlliances"},{"Game::NodeRoute","Game::NodeRoute"},{"Game::FlightPlan::Waypoint","Game::Waypoint"},{"Game::FlightPlan","Game::FlightPlan"},
        {"Game::ShipHealth","Game::ShipHealth"},{"Game::PrisonerHold","Game::PrisonerHold"},{"Game::EventStorage","Game::EventStorage"},{"Game::CivilianRatios","Game::CivilianRatios"},{"Game::ShipRecords","Game::ShipRecords"},
        {"Game::SpyReport","Game::SpyReport"},{"Game::PlayerColorID","Game::PlayerColorID"},{"Mars::Vector3","Mars::Vector3"},
        {"Game::BuildQueue","Game::BuildQueue"},{"Game::TechTree","Game::TechTree"},{"Mars::Stream","Mars::Stream"},{"Game::StrategyEvent","Game::StrategyEvent"}};
    // serializers whose this = object + COL offset : {addr, class type name, col}
    static final Object[][] SHIFTED = {{0x00749630L,"Game::ServerSystem",8},{0x0075d4b0L,"Game::ServerSystem",8},{0x00701070L,"Game::StarFleet",8},{0x00702470L,"Game::StarFleet",8},
        {0x008291f0L,"Game::StarShip",8},{0x00853fa0L,"Game::StarShip",8},{0x008563e0L,"Game::ServerPlayer",0x3a0},{0x008804d0L,"Game::ServerPlayer",0x3a0}};

    DataType findSots(String name) { List<DataType> l = new ArrayList<DataType>(); dtm.findDataTypes(name, l); for (DataType d : l) if (d.getCategoryPath().getPath().equals("/SOTS")) return d; return l.isEmpty() ? null : l.get(0); }
    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;
    }
    void verify(long addr, String outName) throws Exception {
        Function f = getFunctionAt(toAddr(addr));
        DecompileResults res = decomp.decompileFunction(f, 240, monitor);
        PrintWriter w = new PrintWriter(new FileWriter("/tmp/spine/" + outName));
        if (res == null || !res.decompileCompleted()) { w.println("[decompile failed]"); w.close(); return; }
        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); w.println(sb); w.close(); log.println("  verify " + outName + " this-> refs: " + (sb.toString().split("this->").length - 1));
    }

    @Override
    public void run() throws Exception {
        dtm = currentProgram.getDataTypeManager(); st = currentProgram.getSymbolTable(); mem = currentProgram.getMemory();
        log = new PrintWriter(new FileWriter("/tmp/spine/fixthis.log"));
        Map<String, Structure> cls = new HashMap<String, Structure>();
        for (String[] c : CLASSES) {
            String path = c[0], tname = c[1];
            Namespace cur = currentProgram.getGlobalNamespace();
            for (String part : path.split("::")) { Namespace n = st.getNamespace(part, cur); if (n == null) n = st.createClass(cur, part, SourceType.USER_DEFINED); cur = n; }
            if (!(cur instanceof GhidraClass)) { cur = st.convertNamespaceToClass(cur); }
            Structure cs = VariableUtilities.findOrCreateClassStruct((GhidraClass) cur, dtm);
            DataType mine = findSots(tname);
            log.println("class " + path + " -> class struct " + cs.getPathName() + " (len " + cs.getLength() + ")  mine=" + (mine == null ? "null" : mine.getPathName() + " len " + mine.getLength()));
            if (mine == null || mine == cs || mine.isEquivalent(cs) && mine.getPathName().equals(cs.getPathName())) { cls.put(tname, cs); continue; }
            if (mine instanceof Structure) {
                cs.replaceWith(mine); cs.setDescription(mine.getDescription());
                dtm.replaceDataType(mine, cs, false);  // repoint every use of the /SOTS copy to the class struct and drop the copy
                log.println("   merged; class struct now len " + cs.getLength() + " with " + cs.getNumDefinedComponents() + " fields");
            }
            cls.put(tname, cs);
        }
        // shifted serializer views (rebuild from the class structs)
        Map<String, Structure> views = new HashMap<String, Structure>();
        for (Object[] s : SHIFTED) {
            String tname = (String) s[1]; int col = ((Number) s[2]).intValue(); String key = tname + col;
            Structure v = views.get(key);
            if (v == null) {
                Structure src = cls.get(tname);
                StructureDataType nv = new StructureDataType(new CategoryPath("/SOTS"), tname + "_ser" + col, src.getLength() - col, dtm);
                nv.setDescription("Serializer view of " + tname + ": this = object + 0x" + Integer.toHexString(col) + " (IStreamable sub-object); only for " + tname + "::Read/Write. Fields below the sub-object offset are NOT visible here.");
                for (DataTypeComponent c : src.getDefinedComponents()) { if (c.getOffset() < col) continue; try { nv.replaceAtOffset(c.getOffset() - col, c.getDataType(), c.getLength(), c.getFieldName(), c.getComment()); } catch (Exception e) {} }
                v = (Structure) dtm.addDataType(nv, DataTypeConflictHandler.REPLACE_HANDLER); views.put(key, v);
            }
            Function f = getFunctionAt(toAddr((Long) s[0]));
            try {
                f.setCustomVariableStorage(true);
                Parameter p = f.getParameter(0);
                p.setDataType(new PointerDataType(v), SourceType.USER_DEFINED);
                if (f.getParameterCount() > 1) f.getParameter(1).setName("stream", SourceType.USER_DEFINED);
                f.setComment((f.getComment() == null ? "" : f.getComment() + "\n") + "NOTE: this = object + 0x" + Integer.toHexString(col) + " (IStreamable sub-object); typed as " + v.getName() + " (custom storage).");
                log.println("  " + f.getName(true) + ": this -> " + v.getName() + "* (custom storage) param0=" + p.getName() + " storage=" + p.getVariableStorage());
            } catch (Exception e) { log.println("  !! " + f.getName(true) + ": " + e.getMessage()); }
        }
        decomp = new DecompInterface(); decomp.openProgram(currentProgram);
        verify(0x00749630L, "verify_ServerSystem_Write.c"); verify(0x0075d4b0L, "verify_ServerSystem_Read.c"); verify(0x0079fa70L, "verify_StrategyServer_Write.c"); verify(0x00701070L, "verify_StarFleet_Write.c");
        verify(0x008563e0L, "verify_ServerPlayer_Write.c"); verify(0x007d98e0L, "verify_BeginProcessTurn.c"); verify(0x007dc6c0L, "verify_ProcessTurn.c"); verify(0x007598e0L, "verify_ServerSystem_ProcessTurn.c"); verify(0x007d9af0L, "verify_OnPlayerEndTurn.c");
        decomp.dispose(); log.close(); println("fixthis done");
    }
}