/*
 * Copyright 2004-2019 H2 Group. Multiple-Licensed under the MPL 2.0,
 * and the EPL 1.0 (http://h2database.com/html/license.html).
 * Initial Developer: H2 Group
 */
package org.h2.engine;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import org.h2.expression.Expression;
import org.h2.expression.ExpressionColumn;
import org.h2.result.LocalResult;
import org.h2.result.Row;
import org.h2.table.Column;
import org.h2.table.Table;
import org.h2.util.StringUtils;
import org.h2.util.Utils;
import org.h2.value.Value;
import org.h2.value.ValueNull;

Class for gathering and processing of generated keys.
/**
 * Class for gathering and processing of generated keys.
 */
public final class GeneratedKeys {
    
Data for result set with generated keys.
/**
     * Data for result set with generated keys.
     */
    private final ArrayList<Map<Column, Value>> data = Utils.newSmallArrayList();

    
Columns with generated keys in the current row.
/**
     * Columns with generated keys in the current row.
     */
    private final ArrayList<Column> row = Utils.newSmallArrayList();

    
All columns with generated keys.
/**
     * All columns with generated keys.
     */
    private final ArrayList<Column> allColumns = Utils.newSmallArrayList();

    
Request for keys gathering. false if generated keys are not needed, true if generated keys should be configured automatically, int[] to specify column indices to return generated keys from, or String[] to specify column names to return generated keys from. /**
     * Request for keys gathering. {@code false} if generated keys are not needed,
     * {@code true} if generated keys should be configured automatically,
     * {@code int[]} to specify column indices to return generated keys from, or
     * {@code String[]} to specify column names to return generated keys from.
     */
    private Object generatedKeysRequest;

    
Processed table.
/**
     * Processed table.
     */
    private Table table;

    
Remembers columns with generated keys.
Params:
column – 
           table column/**
     * Remembers columns with generated keys.
     *
     * @param column
     *            table column
     */
    public void add(Column column) {
        if (Boolean.FALSE.equals(generatedKeysRequest)) {
            return;
        }
        row.add(column);
    }

    
Clears all information from previous runs and sets a new request for
gathering of generated keys.
Params:
generatedKeysRequest –  false if generated keys are not needed, true if generated keys should be configured automatically, int[] to specify column indices to return generated keys from, or String[] to specify column names to return generated keys from/**
     * Clears all information from previous runs and sets a new request for
     * gathering of generated keys.
     *
     * @param generatedKeysRequest
     *            {@code false} if generated keys are not needed, {@code true} if
     *            generated keys should be configured automatically, {@code int[]}
     *            to specify column indices to return generated keys from, or
     *            {@code String[]} to specify column names to return generated keys
     *            from
     */
    public void clear(Object generatedKeysRequest) {
        this.generatedKeysRequest = generatedKeysRequest;
        data.clear();
        row.clear();
        allColumns.clear();
        table = null;
    }

    
Saves row with generated keys if any.
Params:
tableRow – 
           table row that was inserted/**
     * Saves row with generated keys if any.
     *
     * @param tableRow
     *            table row that was inserted
     */
    public void confirmRow(Row tableRow) {
        if (Boolean.FALSE.equals(generatedKeysRequest)) {
            return;
        }
        int size = row.size();
        if (size > 0) {
            if (size == 1) {
                Column column = row.get(0);
                data.add(Collections.singletonMap(column, tableRow.getValue(column.getColumnId())));
                if (!allColumns.contains(column)) {
                    allColumns.add(column);
                }
            } else {
                HashMap<Column, Value> map = new HashMap<>();
                for (Column column : row) {
                    map.put(column, tableRow.getValue(column.getColumnId()));
                    if (!allColumns.contains(column)) {
                        allColumns.add(column);
                    }
                }
                data.add(map);
            }
            row.clear();
        }
    }

    
Returns generated keys.
Params:
session – 
           session
Returns:
local result with generated keys/**
     * Returns generated keys.
     *
     * @param session
     *            session
     * @return local result with generated keys
     */
    public LocalResult getKeys(Session session) {
        Database db = session.getDatabase();
        if (Boolean.FALSE.equals(generatedKeysRequest)) {
            clear(null);
            return db.getResultFactory().create();
        }
        ArrayList<ExpressionColumn> expressionColumns;
        if (Boolean.TRUE.equals(generatedKeysRequest)) {
            expressionColumns = new ArrayList<>(allColumns.size());
            for (Column column : allColumns) {
                expressionColumns.add(new ExpressionColumn(db, column));
            }
        } else if (generatedKeysRequest instanceof int[]) {
            if (table != null) {
                int[] indices = (int[]) generatedKeysRequest;
                Column[] columns = table.getColumns();
                int cnt = columns.length;
                allColumns.clear();
                expressionColumns = new ArrayList<>(indices.length);
                for (int idx : indices) {
                    if (idx >= 1 && idx <= cnt) {
                        Column column = columns[idx - 1];
                        expressionColumns.add(new ExpressionColumn(db, column));
                        allColumns.add(column);
                    }
                }
            } else {
                clear(null);
                return db.getResultFactory().create();
            }
        } else if (generatedKeysRequest instanceof String[]) {
            if (table != null) {
                String[] names = (String[]) generatedKeysRequest;
                allColumns.clear();
                expressionColumns = new ArrayList<>(names.length);
                for (String name : names) {
                    Column column;
                    search: if (table.doesColumnExist(name)) {
                        column = table.getColumn(name);
                    } else {
                        name = StringUtils.toUpperEnglish(name);
                        if (table.doesColumnExist(name)) {
                            column = table.getColumn(name);
                        } else {
                            for (Column c : table.getColumns()) {
                                if (c.getName().equalsIgnoreCase(name)) {
                                    column = c;
                                    break search;
                                }
                            }
                            continue;
                        }
                    }
                    expressionColumns.add(new ExpressionColumn(db, column));
                    allColumns.add(column);
                }
            } else {
                clear(null);
                return db.getResultFactory().create();
            }
        } else {
            clear(null);
            return db.getResultFactory().create();
        }
        int columnCount = expressionColumns.size();
        if (columnCount == 0) {
            clear(null);
            return db.getResultFactory().create();
        }
        LocalResult result = db.getResultFactory().create(session,
            expressionColumns.toArray(new Expression[0]), columnCount);
        for (Map<Column, Value> map : data) {
            Value[] row = new Value[columnCount];
            for (Map.Entry<Column, Value> entry : map.entrySet()) {
                int idx = allColumns.indexOf(entry.getKey());
                if (idx >= 0) {
                    row[idx] = entry.getValue();
                }
            }
            for (int i = 0; i < columnCount; i++) {
                if (row[i] == null) {
                    row[i] = ValueNull.INSTANCE;
                }
            }
            result.addRow(row);
        }
        clear(null);
        return result;
    }

    
Initializes processing of the specified table. Should be called after clear(), but before other methods. 
Params:
table – 
           table/**
     * Initializes processing of the specified table. Should be called after
     * {@code clear()}, but before other methods.
     *
     * @param table
     *            table
     */
    public void initialize(Table table) {
        this.table = table;
    }

    
Clears unsaved information about previous row, if any. Should be called
before processing of a new row if previous row was not confirmed or simply
always before each row.
/**
     * Clears unsaved information about previous row, if any. Should be called
     * before processing of a new row if previous row was not confirmed or simply
     * always before each row.
     */
    public void nextRow() {
        row.clear();
    }

    @Override
    public String toString() {
        return allColumns + ": " + data.size();
    }

}