intellij-code-from-ds/src/com/rthoni/intellij/codefromds/ui/dialogs/GenerateDialog.java

package com.rthoni.intellij.codefromds.ui.dialogs;

import com.intellij.database.model.DasObject;
import com.intellij.database.model.DasTable;
import com.intellij.database.psi.DbDataSource;
import com.intellij.database.psi.DbPsiFacade;
import com.intellij.database.util.DasUtil;
import com.intellij.openapi.fileChooser.FileChooserDescriptorFactory;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.project.ProjectManager;
import com.intellij.openapi.ui.DialogWrapper;
import com.intellij.openapi.ui.TextFieldWithBrowseButton;
import com.intellij.openapi.ui.ValidationInfo;
import com.intellij.ui.ColoredListCellRenderer;
import com.intellij.ui.JBColor;
import com.intellij.ui.components.JBList;
import com.rthoni.intellij.codefromds.dbo.ColumnSelection;
import com.rthoni.intellij.codefromds.dbo.DataSourceSelection;
import com.rthoni.intellij.codefromds.dbo.GenerateOptions;
import com.rthoni.intellij.codefromds.dbo.TableSelection;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import javax.swing.*;
import javax.swing.event.DocumentEvent;
import javax.swing.event.DocumentListener;
import javax.swing.event.ListSelectionListener;
import java.awt.*;
import java.awt.event.ActionListener;
import java.io.File;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.Vector;
import java.util.stream.Collectors;

/**
 * Created by robin on 11/14/16.
 */
public class GenerateDialog extends DialogWrapper {

    private JPanel _panel;
    private JBList _listDatasources;
    private JBList _listTables;
    private JBList _listColumns;
    private TextFieldWithBrowseButton _textModels;

    private GenerateOptions _options;

    public GenerateDialog(@Nullable Project project) {
        super(project);
        setTitle("Code FROM data source");
        showSource(null);
        init();
    }

    public GenerateOptions getOptions() {
        return _options;
    }

    @Nullable
    @Override
    protected ValidationInfo doValidate() {
        ValidationInfo info = null;
        File modelDir = _options == null ? null : new File(_options.getModelsFolder());
        if (_options == null) {
            info = new ValidationInfo("No Data Source Selected", _listDatasources);
        }
        else if (!modelDir.exists() || !modelDir.isDirectory()) {
            info = new ValidationInfo("Bad Model Folder", _textModels.getTextField());
        }

        return info;
    }

    @Nullable
    @Override
    protected JComponent createCenterPanel() {
        ProjectManager pm = ProjectManager.getInstance();
        Project[] projects = pm.getOpenProjects();

        List<DbDataSource> dataSources = Arrays.stream(projects).map(project -> DbPsiFacade.getInstance(project).getDataSources()).flatMap(Collection::stream).collect(Collectors.toList());

        _listDatasources.setListData(dataSources.stream().map(DasObject::getName).toArray(String[]::new));
        _listDatasources.addListSelectionListener(e -> {
            if (!e.getValueIsAdjusting()) {
                int index = _listDatasources.getSelectedIndex();
                if (index == -1) {
                    showSource(null);
                }
                else {
                    showSource(dataSources.get(index));
                }
            }
        });
        _listTables.addListSelectionListener(e -> {
            if (!e.getValueIsAdjusting()) {
                int index = _listTables.getSelectedIndex();
                if (index == -1) {
                    showTable(null);
                }
                else {
                    showTable(_options.getSelection().getTables().get(index));
                }
            }
        });
        _textModels.getTextField().getDocument().addDocumentListener(new DocumentListener() {
            @Override
            public void insertUpdate(DocumentEvent e) {
                changed();
            }

            @Override
            public void removeUpdate(DocumentEvent e) {
                changed();
            }

            @Override
            public void changedUpdate(DocumentEvent e) {
                changed();
            }

            public void changed() {
                _options.setModelsFolder(_textModels.getText());
            }
        });

        if (dataSources.size() > 0) {
            _listDatasources.setSelectedIndices(new int[]{0});
        }

        return _panel;
    }

    private void showSource(DbDataSource source)
    {
        showTable(null);
        for (ActionListener l : _textModels.getButton().getActionListeners()) {
            _textModels.getButton().removeActionListener(l);
        }
        if (source != null) {
            _options = new GenerateOptions();

            _listTables.setCellRenderer(new ColoredListCellRenderer() {
                @Override
                protected void customizeCellRenderer(@NotNull JList jList, Object o, int i, boolean b, boolean b1) {
                    TableSelection tableSelection = _options.getSelection().getTables().get(i);
                    if (tableSelection.hasAll()) {
                        setBackground(JBColor.GREEN);
                    }
                    else if (tableSelection.hasNone()) {
                        setBackground(JBColor.RED);
                    }
                    else {
                        setBackground(JBColor.ORANGE);
                    }
                    append(tableSelection.getTable().getName());
                }
            });

            _textModels.setText(source.getProject().getBasePath() + File.separator + "Models");
            _textModels.addBrowseFolderListener("Choose Models Destination Folder", "Choose folder",
                    source.getProject(), FileChooserDescriptorFactory.createSingleFolderDescriptor());

            _options.setSelection(new DataSourceSelection());
            _options.getSelection().setSource(source);
            final List<? extends DasTable> tables = DasUtil.getTables(source).toList();
            List<TableSelection> tableSelections = new Vector<>();
            for (DasTable table : tables) {
                TableSelection tableSelection = new TableSelection();
                tableSelection.setTable(table);
                tableSelection.setColumns(DasUtil.getColumns(table).toList().stream().map(c -> {
                    ColumnSelection selection = new ColumnSelection();
                    selection.setColumn(c);
                    selection.setSelected(true);
                    return selection;
                }).collect(Collectors.toList()));
                tableSelections.add(tableSelection);
            }
            _options.getSelection().setTables(tableSelections);

            _listTables.setListData(tables.stream().map(DasObject::getName).toArray(String[]::new));
        }
        else {
            _options = null;
            _listTables.setListData(new String[]{});
        }
    }

    private int[] getSelectedIndices(final TableSelection table)
    {
        List<ColumnSelection> columns = table.getColumns();
        List<Integer> indices = new Vector<>();
        for (int i = 0; i < columns.size(); ++i) {
            if (columns.get(i).isSelected()) {
                indices.add(i);
            }
        }
        return indices.stream().mapToInt(Integer::intValue).toArray();
    }

    private void updateSelection(final TableSelection table)
    {
        List<ColumnSelection> columns = table.getColumns();
        for (int i = 0; i < columns.size(); ++i) {
            columns.get(i).setSelected(_listColumns.isSelectedIndex(i));
        }
    }

    private void showTable(TableSelection table)
    {
        if (table != null) {
            for (ListSelectionListener e : _listColumns.getListSelectionListeners()) {
                _listColumns.removeListSelectionListener(e);
            }
            _listColumns.setListData(table.getColumns().stream().map(c -> c.getColumn().getName()).toArray(String[]::new));
            int[] indices = getSelectedIndices(table);
            _listColumns.setSelectedIndices(indices);
            _listColumns.addListSelectionListener(e -> {
                if (!e.getValueIsAdjusting()) {
                    _listTables.updateUI();
                    updateSelection(table);
                }
            });
        }
        else {
            _listColumns.setListData(new String[]{});
        }
    }
}