// GreenEyeControl - drive a live risc-v-gsp-rm-debugger session from Ghidra.
//
// The script connects to the risc-v-gsp-rm-debugger JSON-lines bridge (start it with
// `./risc-v-gsp-rm-debugger.sh ghidra [--kernel <elf>]`), and exposes the operations a
// reverse engineer needs while reading the decompiler:
//
//   ping()               - check the bridge
//   info()               - session name, pc, symbol count, recording path
//   resolve(name)        - symbol/expression -> runtime address
//   setBreakpointHere()  - software breakpoint at currentAddress
//   setBreakpointAt(addr) / clearBreakpointAt(addr)
//   runUntilHere()       - resume and stop when execution reaches currentAddress
//   runUntil(addr)       - same for an explicit address (hex string)
//   readMemoryHere(size) - hex bytes at currentAddress
//   readRegister(name) / writeRegister(name, value)
//   callFunction(addr, args) - Frida-like remote call; returns a0
//   snapshotHere(size)   - record a memory snapshot in the session
//   listHits()           - breakpoint hit counts from the session
//   exportHits(path)     - write greeneye_hits.tsv for MarkGreenEyeHits.java
//
// Usage (Ghidra headless):
//   analyzeHeadless <proj> <name> -process <program> -noanalysis \
//       -scriptPath debug/risc-v-gsp-rm-vm/ghidra \
//       -postScript GreenEyeControl.java <bridge-port> [address]
//
// The optional address argument makes the script run until that address and
// print the hits; in the GUI, omit it to use the current cursor address.
//
//@category OpenIOV
//@menupath Tools.OpenIOV.GreenEyeControl

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.PrintWriter;
import java.net.Socket;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;

public class GreenEyeControl extends GhidraScript {

    private Socket socket;
    private BufferedReader reader;
    private PrintWriter writer;
    private int requestId = 0;

    private static final Pattern OK = Pattern.compile("\"ok\"\\s*:\\s*(true|false)");
    private static final Pattern RESULT = Pattern.compile("\"result\"\\s*:\\s*(\\{.*\\}|\\[.*\\]|\"[^\"]*\"|[^,}]+)");

    @Override
    public void run() throws Exception {
        String[] args = getScriptArgs();
        if (args.length < 1) {
            println("usage: GreenEyeControl.java <bridge-port> [address]");
            return;
        }
        int port = Integer.parseInt(args[0]);
        try {
            connect(port);
            println("risc-v-gsp-rm-debugger: " + info());
            String target = args.length > 1 ? args[1] : null;
            if (target == null && currentAddress != null) {
                target = "0x" + Long.toHexString(currentAddress.getOffset());
            }
            if (target != null) {
                println("risc-v-gsp-rm-debugger run_until " + target + " -> " + runUntil(target));
                println("risc-v-gsp-rm-debugger hits: " + listHits());
            }
        } finally {
            disconnect();
        }
    }

    // -- connection --------------------------------------------------------

    public void connect(int port) throws Exception {
        socket = new Socket("127.0.0.1", port);
        reader = new BufferedReader(new InputStreamReader(socket.getInputStream()));
        writer = new PrintWriter(socket.getOutputStream(), true);
    }

    public void disconnect() throws Exception {
        if (socket != null) {
            socket.close();
        }
    }

    /** Send one bridge command and return the raw response line. */
    public String send(String cmd, String argsJson) throws Exception {
        requestId += 1;
        String request = "{\"id\":" + requestId + ",\"cmd\":\"" + cmd + "\",\"args\":"
            + (argsJson == null || argsJson.isEmpty() ? "{}" : argsJson) + "}";
        writer.println(request);
        String response = reader.readLine();
        if (response == null) {
            throw new Exception("risc-v-gsp-rm-debugger bridge closed the connection");
        }
        Matcher ok = OK.matcher(response);
        if (ok.find() && "false".equals(ok.group(1))) {
            throw new Exception("risc-v-gsp-rm-debugger error: " + response);
        }
        return response;
    }

    /** Extract the `result` field as a string (numbers, objects or lists). */
    public String resultOf(String response) {
        Matcher matcher = RESULT.matcher(response);
        return matcher.find() ? matcher.group(1) : response;
    }

    // -- documented operations --------------------------------------------

    public String ping() throws Exception {
        return resultOf(send("ping", null));
    }

    public String info() throws Exception {
        return resultOf(send("info", null));
    }

    public String resolve(String name) throws Exception {
        return resultOf(send("resolve", "{\"name\":\"" + name + "\"}"));
    }

    public String setBreakpointAt(String address) throws Exception {
        return resultOf(send("bp_set", "{\"target\":\"" + address + "\"}"));
    }

    public String setBreakpointHere() throws Exception {
        return setBreakpointAt(addressHex(currentAddress));
    }

    public String clearBreakpointAt(long address) throws Exception {
        return resultOf(send("bp_clear", "{\"addr\":" + address + "}"));
    }

    public String runUntil(String address) throws Exception {
        return resultOf(send("run_until", "{\"target\":\"" + address + "\",\"timeout\":10}"));
    }

    public String runUntilHere() throws Exception {
        return runUntil(addressHex(currentAddress));
    }

    public String readMemoryHere(int size) throws Exception {
        return resultOf(send("read_mem", "{\"addr\":\"" + addressHex(currentAddress) + "\",\"size\":" + size + "}"));
    }

    public String readRegister(String name) throws Exception {
        return resultOf(send("read_reg", "{\"name\":\"" + name + "\"}"));
    }

    public String writeRegister(String name, long value) throws Exception {
        return resultOf(send("write_reg", "{\"name\":\"" + name + "\",\"value\":" + value + "}"));
    }

    public String callFunction(String address, long... args) throws Exception {
        StringBuilder json = new StringBuilder("{\"target\":\"" + address + "\",\"args\":[");
        for (int i = 0; i < args.length; i++) {
            if (i > 0) {
                json.append(",");
            }
            json.append(args[i]);
        }
        json.append("]}");
        return resultOf(send("call", json.toString()));
    }

    public String snapshotHere(int size) throws Exception {
        return resultOf(send("snapshot",
            "{\"addr\":\"" + addressHex(currentAddress) + "\",\"size\":" + size + ",\"label\":\"ghidra\"}"));
    }

    public String listHits() throws Exception {
        return resultOf(send("hits", null));
    }

    public String exportHits(String path) throws Exception {
        return resultOf(send("export_hits", "{\"path\":\"" + path + "\"}"));
    }

    private String addressHex(Address address) {
        return "0x" + Long.toHexString(address.getOffset());
    }
}