// Imports data from a file produced by the EnhancedExport script.
//
// @category Import

import ghidra.app.script.GhidraScript;
import ghidra.app.services.DataTypeQueryService;
import ghidra.program.model.address.Address;
import ghidra.program.model.data.ArrayDataType;
import ghidra.program.model.data.DataType;
import ghidra.program.model.data.PointerDataType;
import ghidra.program.model.data.Undefined4DataType;
import ghidra.program.model.listing.Data;
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.symbol.Namespace;
import ghidra.program.model.symbol.SourceType;
import ghidra.program.model.symbol.Symbol;
import ghidra.util.StringUtilities;

import java.io.FileReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Optional;


public class EnhancedImport extends GhidraScript{
	@Override
	public void run() throws Exception {
		final FileReader   in    = new FileReader(askFile("Select input file", "OK"));
		final List<String> lines = in.readAllLines();
		in.close();

		for (int i = 0; i < lines.size(); i++) {
			final String[] parts = lines.get(i).split("\t");;
			final Address  addr  = toAddr(parts[0]);

			switch (parts[1]) {
				case "data": importData(addr, parts); break;
				case "func": importFunc(addr, parts); break;
				default: throw new Exception(
					"Symbol type \"" + parts[1] +
					"\" on line " + String.valueOf(i) +
					" is not \"data\" or \"func\""
				);
			}
		}
	}

	private void importData(
		final Address  addr,
		final String[] parts
	) throws Exception {
		print("Importing data symbol \"" + parts[3] + "\"...");

		// Create symbol
		final Namespace ns = importNamespace(parts[3]);
		final Symbol    s  = Optional.ofNullable(getSymbolAt(addr))
		                             .orElse(createLabel(
		                             	addr,
		                             	unqualified(parts[3]),
		                             	ns,
		                             	true,
		                             	SourceType.USER_DEFINED
		                             ));

		// Create data
		if (makeType(parts[2]).orElse(null) instanceof DataType t) {
			clearListing(addr, addr.add(Math.max(t.getLength(), 1) - 1));
			currentProgram.getListing().createData(addr, t);

			println(" done.");
		} else println(" skipping.");
	}

	private static String unqualified(final String qualifiedName) {
	/* Strips the namespaces off of a qualified name */
		final String[] parts = qualifiedName.split("::");
		return parts[parts.length - 1];
	}

	private Namespace importNamespace(final String qualifiedName) throws Exception {
	/* Creates namespaces from the given name, returning the deepest one
	Returns null if the qualified name is in the global namespace.
	*/
		final String[] parts = qualifiedName.split("::");

		if (parts.length < 2) return null;

		final String[] names = Arrays.copyOfRange(parts, 0, parts.length - 1);

		Namespace ns = null;
		for (final String name : names) ns = createNamespace(ns, name);

		return ns;
	}

	private Optional<DataType> makeType(final String type) throws Exception {
	/* Attempt to create the described type from a known base type */
		final String         baseName   = StringUtilities.findWord(type, 0);
		final List<DataType> foundTypes = state.getTool()
		                                       .getService(DataTypeQueryService.class)
		                                       .findDataTypes(baseName, null);
		if (foundTypes.size() == 0) {
			print(" can't find data type \"" + baseName + "\",");
			return Optional.empty();
		}

		return Optional.of(derivedType(
			foundTypes.get(0), // Boldly assume first is right
			type.substring(baseName.length())
		));
	}

	private static DataType derivedType(
		final DataType t,
		final String   modifiers
	) throws Exception {
	/* Return the given datatype with given pointer/array modifiers */
		DataType ret = t;
		for (int i = 0; i < modifiers.length(); i++) switch (modifiers.charAt(i)) {
			case '*':
				ret = new PointerDataType(ret);
				continue;

			case '[':
				// Read array dimensions
				final ArrayList<Integer> dims = new ArrayList<>();
				do {
					final String n  = StringUtilities.findWord(modifiers, i+1);
					i              += n.length() + 1;

					dims.add(Integer.valueOf(n));
				} while (modifiers.charAt(i) == '[');

				// Apply array
				Collections.reverse(dims);
				for (int dim : dims) ret = new ArrayDataType(ret, dim);

				continue;

			case ' ':
				continue;

			default:
				throw new Exception(
					"Unexpected character \""           +
					String.valueOf(modifiers.charAt(i)) +
					"\" in type modifiers"
				);
		}

		return ret;
	}

	private void importFunc(
		final Address  addr,
		final String[] parts
	) throws Exception {
		print("Importing function symbol \"" + parts[6] + "\"...");

		final Function f = Optional.ofNullable(getFunctionAt(addr))
		                           .orElse(createFunction(addr, parts[6]));

		if (makeType(parts[2]).orElse(null) instanceof DataType t)
			f.setReturnType(t, SourceType.USER_DEFINED);
		f.setInline(parts[4].equals("inline"));
		f.setCallFixup(parts[5].equals("nofixup") ? null : parts[5]);
		f.setName(unqualified(parts[6]), SourceType.USER_DEFINED);
		if (importNamespace(parts[6]) instanceof Namespace ns)
			f.setParentNamespace(ns);

		final ArrayList<Parameter> args = new ArrayList<>();
		for (int i = 7; i < parts.length - 1; i += 2)
			args.add(new ParameterImpl(
				parts[i+1],
				makeType(parts[i]).orElse(Undefined4DataType.dataType),
				currentProgram
			));

		f.updateFunction(
			parts[3],
			null,
			args,
			FunctionUpdateType.DYNAMIC_STORAGE_FORMAL_PARAMS,
			true,
			SourceType.USER_DEFINED
		);

		f.setVarArgs(parts[parts.length - 1].equals("..."));

		println(" done.");
	}
}