import java.awt.BorderLayout;
import java.io.File;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.regex.Pattern;

import javax.swing.JDialog;
import javax.swing.JFileChooser;
import javax.swing.JOptionPane;
import javax.swing.JPanel;
import javax.swing.JScrollPane;
import javax.swing.JTable;
import javax.swing.JTextArea;
import javax.swing.table.DefaultTableModel;
import javax.swing.table.TableColumn;
import javax.swing.table.TableRowSorter;

import ghidra.app.script.GhidraScript;
import ghidra.app.util.datatype.DataTypeSelectionDialog;
import ghidra.framework.plugintool.PluginTool;
import ghidra.program.model.data.DataType;
import ghidra.program.model.data.DataTypeManager;
import ghidra.program.model.lang.CompilerSpec;
import ghidra.program.model.lang.PrototypeModel;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.FunctionIterator;
import ghidra.program.model.listing.FunctionTag;
import ghidra.util.data.DataTypeParser.AllowedDataTypes;
import ghidra.util.exception.CancelledException;

public class BulkUtils {
	
	private static final Set<String> RESERVED_C_KEYWORDS = Set.of(
            "auto", "break", "case", "char", "const", "continue", "default", "do", "double", "else",
            "enum", "extern", "float", "for", "goto", "if", "inline", "int", "long", "register",
            "restrict", "return", "short", "signed", "sizeof", "static", "struct", "switch",
            "typedef", "union", "unsigned", "void", "volatile", "while", "_Alignas", "_Alignof",
            "_Atomic", "_Bool", "_Complex", "_Generic", "_Imaginary", "_Noreturn", "_Static_assert",
            "_Thread_local"
    );
	private static final Pattern C_PARAM_REGEX = Pattern.compile("[a-zA-Z_][a-zA-Z\\\\d_]*"); 
	
	static File askForOutputFile(String suggestedFileName) {
		JFileChooser fileChooser = new JFileChooser();
		if(suggestedFileName != null) {
			fileChooser.setSelectedFile(new File(suggestedFileName));			
		}
        int fileChooserResult = fileChooser.showSaveDialog(null);
        if(fileChooserResult == JFileChooser.APPROVE_OPTION) {
        	File outputFile = fileChooser.getSelectedFile();
        	if(outputFile.exists()) {
        		String question = "This file already exists, do you want to overwrite it?";
        		int confirmOverwriteResult = JOptionPane.showConfirmDialog(null, question, null, JOptionPane.YES_NO_OPTION);
        		if(confirmOverwriteResult != JOptionPane.YES_OPTION) {
        			return null;
        		}
        	}
        	return outputFile;
        }
        return null;
	}
	
	static boolean isValidCParameterName(String parameterName) {
		if(parameterName == null || parameterName.isEmpty()) return false;
		if(RESERVED_C_KEYWORDS.contains(parameterName)) return false;
		return C_PARAM_REGEX.matcher(parameterName).matches();
	}
	
	static DataType askForDataType(GhidraScript gs, String title) {
		DataTypeManager dataTypeManager = gs.getCurrentProgram().getDataTypeManager();
        PluginTool tool = gs.getState().getTool();
        DataTypeSelectionDialog typeDialog = new DataTypeSelectionDialog(tool, dataTypeManager, -1, AllowedDataTypes.FIXED_LENGTH);
        if(title != null) {
        	typeDialog.setTitle(title);
        }
        tool.showDialog(typeDialog);
        return typeDialog.getUserChosenDataType();
	}
	
	static void printDoneMessage(GhidraScript gs, int success, int fail) {
		gs.println("Done: successful: " + success + " | failed: " + fail);
	}
	
	static String askForCallingConvention(GhidraScript gs) throws CancelledException {
		CompilerSpec compilerSpec = gs.getCurrentProgram().getCompilerSpec();
		PrototypeModel[] callingConventions = compilerSpec.getCallingConventions();
		List<String> choices = new ArrayList<>();
		choices.add(Function.DEFAULT_CALLING_CONVENTION_STRING);
		choices.add(Function.UNKNOWN_CALLING_CONVENTION_STRING);
		for(PrototypeModel callingConvention : callingConventions) {
			choices.add(callingConvention.getName());
		}
		return gs.askChoice(null, "Calling convention:", choices, choices.get(0));
	}
	
	static Integer askForInteger(GhidraScript gs, String whatFor) throws CancelledException {
		String input = gs.askString(null, whatFor);
		try {
			return Integer.parseInt(input, 10);
		} catch(NumberFormatException nfe) {}
		return null;
	}
	
	static Pattern askForRegex(GhidraScript gs, String whatToMatch) throws CancelledException {
		String regex = gs.askString("Enter Regex", "Pattern to match " + whatToMatch + ":");
		Pattern pattern = Pattern.compile(regex);
		return pattern;
	}
	
	static List<Function> getFunctions(GhidraScript gs) throws CancelledException{
		List<String> modeChoices = Arrays.asList("Name Regex", "Tags");
		String modeSelection = gs.askChoice("Selection Mode", "How do you want to select functions?", modeChoices, modeChoices.get(0));
		switch(modeSelection) {
		default:
		case "Name Regex":
			Pattern pattern = askForRegex(gs, "entire function name");
			return getFunctionsByRegex(gs, pattern);
		case "Tags":
			String[] tags = gs.askString("Tags", "select for (comma separated):").split(",");
			String tagModeSelection;
			switch(tags.length) {
			case 0:
				System.out.println("Abort: no tags provided");
				throw new CancelledException();
			case 1:
				tagModeSelection = "AND";
				break;
			default:
				List<String> tagModeChoices = Arrays.asList("AND", "OR");
				tagModeSelection = gs.askChoice("Tag Mode", "Behavior for multiple tags", tagModeChoices, tagModeChoices.get(0));
				break;
			}
			return getFunctionsByTags(gs, tagModeSelection == "AND", tags);
		}
	}
	
	static List<Function> getFunctionsByRegex(GhidraScript gs, Pattern pattern) {
		List<Function> functions = new ArrayList<>();
		FunctionIterator iterator = gs.getCurrentProgram().getFunctionManager().getFunctions(true);
		while (iterator.hasNext() && !gs.getMonitor().isCancelled()) {
            Function function = iterator.next();
            if (pattern.matcher(function.getName()).matches()) {
            	functions.add(function);
            }
        }
		return functions;
	}
	
	static List<Function> getFunctionsByTags(GhidraScript gs, boolean useAND, String... searchTags){
		List<String> searchTagsList = Arrays.asList(searchTags);
		List<Function> functions = new ArrayList<>();
		FunctionIterator iterator = gs.getCurrentProgram().getFunctionManager().getFunctions(true);
		boolean matches;
		while (iterator.hasNext() && !gs.getMonitor().isCancelled()) {
            Function function = iterator.next();
            Set<FunctionTag> functionTags = function.getTags();
            if(functionTags.size() == 0) continue; // functions without tags are irrelevant
            if(useAND && functionTags.size() < searchTags.length) continue; // search mode is AND + function has fewer tags as we are searching for
            matches = useAND ? true : false;
            boolean tagContained;
            for(FunctionTag functionTag : functionTags)
            {
            	tagContained = searchTagsList.contains(functionTag.getName());
            	if(useAND)
            	{
            		 if(!tagContained) {
            			 matches = false;
            			 break;
            		 }
            	} else if(tagContained) {
            		matches = true;
            		break;
            	}
            }
            if(matches){
            	functions.add(function);
            }
        }
		return functions;
	}
	
	static int removeFunctionsByCallingConvention(List<Function> functions, String callingConvention) {
		int cnt = 0;
		Iterator<Function> iterator = functions.iterator();
		while(iterator.hasNext()) {
			Function function = iterator.next();
			if(function.getCallingConventionName().equals(callingConvention)) {
				iterator.remove();
				cnt++;
			}
		}
		return cnt;
	}
	
	static int removeFunctionsByDataType(List<Function> functions, DataType dataType) {
		int cnt = 0;
		Iterator<Function> iterator = functions.iterator();
		while(iterator.hasNext()) {
			Function function = iterator.next();
			if(function.getReturnType().equals(dataType)) {
				iterator.remove();
				cnt++;
			}
		}
		return cnt;
	}
	
	static boolean showFunctionConfirmationDialog(List<Function> functions, String title, int skipped, String skippedString) {
		DefaultTableModel tableModel = new DefaultTableModel(new String[] {"Name", "Entry Point"}, 0);
		functions.forEach(function -> {
			tableModel.addRow(new String[] {function.getName(), function.getEntryPoint().toString().toUpperCase()});
		});
        TableRowSorter<DefaultTableModel> sorter = new TableRowSorter<>(tableModel);
        sorter.setSortsOnUpdates(true);
        sorter.toggleSortOrder(0); // Sort by first column (Name)
        JTable table = new JTable(tableModel);
        table.setRowSorter(sorter);
		TableColumn entryPointColumn = table.getColumnModel().getColumn(1);
		int fixedWidthEntryPoint = 130;
        entryPointColumn.setMinWidth(fixedWidthEntryPoint);
        entryPointColumn.setMaxWidth(fixedWidthEntryPoint);
        entryPointColumn.setPreferredWidth(fixedWidthEntryPoint);
        JScrollPane scrollPane = new JScrollPane(table);
        JPanel panel = new JPanel(new BorderLayout());
        if(skippedString != null) {
        	JTextArea skippedText = new JTextArea(String.format(skippedString, skipped));
    		skippedText.setEditable(false);
    		panel.add(skippedText, BorderLayout.NORTH);
        }
        panel.add(scrollPane, BorderLayout.CENTER);
		Object[] options = { "OK", "Cancel" };
		JOptionPane optionPane = new JOptionPane(panel, JOptionPane.PLAIN_MESSAGE, JOptionPane.OK_CANCEL_OPTION, null, options, options[1]); // select cancel initially
        JDialog confirmDialog = optionPane.createDialog(null, title);
		confirmDialog.setResizable(true);
        confirmDialog.setVisible(true);
		return optionPane.getValue() == options[0];
	}
	
}