// Apply Wine-11.15-backed Win16 import prototypes before decompilation.
// @category TDKPIN

import java.io.BufferedReader;
import java.io.FileReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import ghidra.app.script.GhidraScript;
import ghidra.program.model.data.ByteDataType;
import ghidra.program.model.data.CharDataType;
import ghidra.program.model.data.DataType;
import ghidra.program.model.data.IntegerDataType;
import ghidra.program.model.data.PointerDataType;
import ghidra.program.model.data.ShortDataType;
import ghidra.program.model.data.UnsignedIntegerDataType;
import ghidra.program.model.data.UnsignedShortDataType;
import ghidra.program.model.data.VoidDataType;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.Function.FunctionUpdateType;
import ghidra.program.model.listing.Parameter;
import ghidra.program.model.listing.ParameterImpl;
import ghidra.program.model.listing.ReturnParameterImpl;
import ghidra.program.model.symbol.SourceType;

public class ApplyWin16ImportSignatures extends GhidraScript {
    private static final Pattern PROTOTYPE = Pattern.compile(
        "^pascal16(?<flags>(?: -\\S+)*) [A-Za-z0-9_@]+\\((?<arguments>.*)\\)$");

    @Override
    public void run() throws Exception {
        String[] args = getScriptArgs();
        if (args.length != 1) {
            throw new IllegalArgumentException(
                "usage: ApplyWin16ImportSignatures.java IMPORT_RECONSTRUCTION.tsv");
        }

        int applied = 0;
        int dataExports = 0;
        int registerExports = 0;
        try (BufferedReader reader = new BufferedReader(new FileReader(args[0]))) {
            String header = reader.readLine();
            if (header == null) {
                throw new IllegalArgumentException("empty import ledger");
            }
            Map<String, Integer> columns = columns(header);
            String line;
            while ((line = reader.readLine()) != null && !monitor.isCancelled()) {
                String[] fields = line.split("\\t", -1);
                String status = field(fields, columns, "status");
                if ("runtime-data".equals(status)) {
                    dataExports++;
                    continue;
                }
                String address = field(fields, columns, "address");
                String prototype = field(fields, columns, "prototype");
                Function function = getFunctionAt(toAddr(address));
                if (function == null) {
                    throw new IllegalStateException("missing imported function at " + address);
                }
                if (prototype.startsWith("pascal16 -register ")) {
                    registerExports++;
                    continue;
                }
                apply(function, prototype);
                applied++;
            }
        }
        analyzeAll(currentProgram);
        println("Applied " + applied + " Win16 import signatures; skipped " + dataExports +
            " runtime-data and " + registerExports + " register-entry slots");
    }

    private void apply(Function function, String prototype) throws Exception {
        Matcher matcher = PROTOTYPE.matcher(prototype);
        if (!matcher.matches()) {
            throw new IllegalArgumentException(
                "unsupported prototype for " + function.getEntryPoint() + ": " + prototype);
        }

        List<Parameter> parameters = new ArrayList<>();
        String arguments = matcher.group("arguments").trim();
        if (!arguments.isEmpty()) {
            String[] tokens = arguments.contains(",") ? arguments.split(",") : arguments.split("\\s+");
            for (int index = 0; index < tokens.length; index++) {
                String token = tokens[index].trim();
                String[] words = token.split("\\s+");
                DataType type = type(words[0]);
                String name = words.length > 1 ? words[words.length - 1] : "arg" + (index + 1);
                parameters.add(new ParameterImpl(name, type, currentProgram));
            }
        }

        boolean ret16 = matcher.group("flags").contains("-ret16");
        DataType returnType = ret16 ? UnsignedShortDataType.dataType
            : UnsignedIntegerDataType.dataType;
        ReturnParameterImpl returnParameter = new ReturnParameterImpl(returnType, currentProgram);
        function.updateFunction("__stdcall16far", returnParameter,
            FunctionUpdateType.DYNAMIC_STORAGE_ALL_PARAMS, true, SourceType.USER_DEFINED,
            parameters.toArray(Parameter[]::new));
    }

    private DataType type(String name) {
        return switch (name) {
            case "byte" -> ByteDataType.dataType;
            case "s_byte" -> ByteDataType.dataType;
            case "word", "HWND16" -> UnsignedShortDataType.dataType;
            case "s_word" -> ShortDataType.dataType;
            case "long" -> UnsignedIntegerDataType.dataType;
            case "s_long" -> IntegerDataType.dataType;
            case "ptr", "segptr" ->
                new PointerDataType(VoidDataType.dataType, 4, currentProgram.getDataTypeManager());
            case "str", "segstr" ->
                new PointerDataType(CharDataType.dataType, 4, currentProgram.getDataTypeManager());
            default -> throw new IllegalArgumentException("unsupported Win16 argument type: " + name);
        };
    }

    private Map<String, Integer> columns(String header) {
        Map<String, Integer> result = new HashMap<>();
        String[] names = header.split("\\t", -1);
        for (int index = 0; index < names.length; index++) {
            result.put(names[index], index);
        }
        return result;
    }

    private String field(String[] fields, Map<String, Integer> columns, String name) {
        Integer index = columns.get(name);
        if (index == null || index >= fields.length) {
            throw new IllegalArgumentException("missing import-ledger column: " + name);
        }
        return fields[index];
    }
}