/*
 * JBoss, Home of Professional Open Source.
 * See the COPYRIGHT.txt file distributed with this work for information
 * regarding copyright ownership.  Some portions may be licensed
 * to Red Hat, Inc. under one or more contributor license agreements.
 * 
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 * 
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 * 
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 * 02110-1301 USA.
 * Modified by
 * @author  Data Virtuality GmbH
 * @date    $Date$
 *  
 */
/*
 * EXTENSION
 * ISSUE: DVCORE-2521
 * FIXED BY: Roman Pedchenko, Salvatore Raunich
 * DATE: 2014-12-05, 2014-12-16
 * ADOPTED FROM: TEIID-2153, 2171, 2359, 2149, 2358, 2410, 2271, 2410, 2467, 2429, 2363, 2160, 2461, 2797, 2326, 1557, 2576
 * END EXTENSION
 */
// https://issues.jboss.org/browse/TEIID-3271
// https://issues.jboss.org/browse/TEIID-3523
package org.teiid.translator.jdbc;

import com.datavirtuality.dv.core.interfaces.DVCaller;
import org.teiid.GeneratedKeys;
import org.teiid.dqp.internal.process.DQPWorkContext;
import org.teiid.language.BatchedCommand;
import org.teiid.language.BatchedUpdates;
import org.teiid.language.Command;
import org.teiid.language.Insert;
import org.teiid.logging.LogConstants;
import org.teiid.logging.LogManager;
import org.teiid.metadata.Column;
import org.teiid.translator.*;
import org.teiid.translator.jdbc.plugin.JDBCPlugin;

import java.sql.*;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;

/**
 */
public class JDBCUpdateExecution extends JDBCBaseExecution implements UpdateExecution {

	// EXTENSION
	// ISSUE: DVCORE-1115
	// FIXED BY: Salvatore Raunich
	// DATE: 2014.09.09
	//	private int[] result;
	protected int[] result;
	// END EXTENSION
	
    /**
     * @param connection
     * @param sqlTranslator
     * @param logger
     * @param props
     * @param id
     */
	public JDBCUpdateExecution(Command command, Connection connection, ExecutionContext context, JDBCExecutionFactory env) {
        super(command, connection, context, env);
    }

    // ===========================================================================================================================
    // Methods
    // ===========================================================================================================================

    @Override
    public void execute() throws TranslatorException {
		Command command1 = preTranslateUpdate(command);

		if (command1 != null)
			command = command1;

		if (command instanceof BatchedUpdates) {
			execute(((BatchedUpdates) command));
		} else {
            // translate command
            TranslatedCommand translatedComm = translateCommand(command);

            try {
    	    	DVCaller.callWithExc(76, command, connection);
    		} catch (Throwable e) {
    			throw new TranslatorException(e);
    		}
            try {
            	executeTranslatedCommand(translatedComm);
    		} catch (TranslatorException e) {
    			DVCaller.call(91, command);
    			throw e;
    		}
        }
    }

	public int[] execute(BatchedUpdates batchedCommand) throws TranslatorException {
        boolean succeeded = false;

        boolean commitType = getAutoCommit(null);
        Command[] commands = batchedCommand.getUpdateCommands().toArray(new Command[batchedCommand.getUpdateCommands().size()]);
        result = new int[commands.length];

        TranslatedCommand tCommand = null;
    	// EXTENSION
    	// ISSUE: DVCORE-1296
    	// FIXED BY: Salvatore Raunich
    	// DATE: 2014.06.02
        JDBCExecutionException exc = null;
        // END EXTENSION
        try {
            // temporarily turn the auto commit off, and set it back to what it was
            // before at the end of the command execution.
            if (commitType) {
                connection.setAutoCommit(false);
            }

            List<TranslatedCommand> executedCmds = new ArrayList<TranslatedCommand>();

            TranslatedCommand previousCommand = null;
            
            for (int i = 0; i < commands.length; i++) {
            	tCommand = translateCommand(commands[i]);
                if (tCommand.isPrepared()) {
                    PreparedStatement pstmt = null;
                    if (previousCommand != null && previousCommand.isPrepared() && previousCommand.getSql().equals(tCommand.getSql())) {
                        pstmt = (PreparedStatement)statement;
                    } else {
                        if (!executedCmds.isEmpty()) {
                            executeBatch(i, result, executedCmds);
                        }
                        pstmt = getPreparedStatement(tCommand.getSql());
                    }
                    bind(pstmt, tCommand.getPreparedValues(), null);
                    pstmt.addBatch();
                } else {
                    if (previousCommand != null && previousCommand.isPrepared()) {
                        executeBatch(i, result, executedCmds);
                        getStatement();
                    }
                    if (statement == null) {
                        getStatement();
                    }
                    statement.addBatch(tCommand.getSql());
                }
                executedCmds.add(tCommand);
                previousCommand = tCommand;
            }
            if (!executedCmds.isEmpty()) {
                executeBatch(commands.length, result, executedCmds);
            }
            succeeded = true;
        } catch (SQLException e) {
        	// EXTENSION
        	// ISSUE: DVCORE-1296
        	// FIXED BY: Salvatore Raunich
        	// DATE: 2014.06.02
//        	throw new JDBCExecutionException(JDBCPlugin.Event.TEIID11011, e, tCommand);
        	exc = new JDBCExecutionException(JDBCPlugin.Event.TEIID11011, e, tCommand);
        	throw exc;
        	// END EXTENSION
        } finally {
            if (commitType) {
            	// EXTENSION
            	// ISSUE: DVCORE-1296
            	// FIXED BY: Salvatore Raunich
            	// DATE: 2014.06.02
//            	restoreAutoCommit(!succeeded, null);
                restoreAutoCommit(!succeeded, null, exc);
                // END EXTENSION
            }
        }

        return result;
    }

    private void executeBatch(int commandCount,
                              int[] results,
                              List<TranslatedCommand> commands) throws TranslatorException {
        try {
            int[] batchResults = statement.executeBatch();
            addStatementWarnings();
            int length = Math.min(batchResults.length, results.length);
			for (int j = 0; j < length; j++) {
                int i = length - 1 - j;
                int result = 0;
                if (i >= 0 && i < length)
                	result = batchResults[i];
				results[length - 1 - j] = result;
            }
            commands.clear();
        } catch (SQLException err) {
             throw new JDBCExecutionException(JDBCPlugin.Event.TEIID11012, err, commands.toArray(new TranslatedCommand[commands.size()]));
        }
    }

    /**
     * @param translatedComm
     * @throws TranslatorException
     * @since 4.3
     */
    protected void executeTranslatedCommand(TranslatedCommand translatedComm) throws TranslatorException {
        // create statement or PreparedStatement and execute
        String sql = translatedComm.getSql();
		DQPWorkContext.getWorkContext().setTranslatedSql(sql);
        boolean commitType = false;
        boolean succeeded = false;
    	// EXTENSION
    	// ISSUE: DVCORE-1296
    	// FIXED BY: Salvatore Raunich
    	// DATE: 2014.06.02
        JDBCExecutionException exc = null;
        // END EXTENSION
        try {
        	int updateCount = 0;
	        String[] keyColumnNames = null;
	        if (command instanceof Insert && context.getCommandContext().isReturnAutoGeneratedKeys() && executionFactory.supportsGeneratedKeys(context, command)) {
		        List<Column> cols = context.getGeneratedKeyColumns();
		        if (cols != null) {
			        //TODO won't work in scenarios where the teiid name is changed or contains a .
			        keyColumnNames = context.returnGeneratedKeys().getColumnNames();
		        }
	        }
        	if (!translatedComm.isPrepared()) {
		        //handle autoGeneratedKeys
		        if (keyColumnNames != null) {
			        if (executionFactory.useColumnNamesForGeneratedKeys()) {
				        updateCount = statement.executeUpdate(sql, keyColumnNames);
			        } else {
				        updateCount = statement.executeUpdate(sql, Statement.RETURN_GENERATED_KEYS);
			        }
		        } else {
        			updateCount = getStatement().executeUpdate(sql);
        		}
		        result = new int[] {updateCount};
        		addStatementWarnings();
        	} else {
		        PreparedStatement pstatement = null;
		        if (statement != null) {
			        statement.close();
		        }
		        if (keyColumnNames != null) {
			        if (executionFactory.useColumnNamesForGeneratedKeys()) {
				        pstatement = connection.prepareStatement(sql, keyColumnNames);
			        } else {
				        pstatement = connection.prepareStatement(sql, Statement.RETURN_GENERATED_KEYS);
			        }
		        } else {
			        pstatement = getPreparedStatement(sql);
		        }
        		Iterator<? extends List<?>> vi = null;
        		if (command instanceof BatchedCommand) {
        			BatchedCommand batchCommand = (BatchedCommand)command;
        			vi = batchCommand.getParameterValues();
        		}
        		if (vi != null) {
        			commitType = getAutoCommit(translatedComm);
        			if (commitType) {
        				connection.setAutoCommit(false);
        			}
        			int maxBatchSize = (command instanceof Insert)?this.executionFactory.getMaxPreparedInsertBatchSize():Integer.MAX_VALUE;
        			boolean done = false;
        			outer: while (!done) {
        				for (int i = 0; i < maxBatchSize; i++) {
        					if (vi.hasNext()) {
        						List<?> values = vi.next();
        						bind(pstatement, translatedComm.getPreparedValues(), values);
        					} else {
        						if (i == 0) {
        							// otherwise we will get NPE when inserting data (select * into)
							        // from another data source with limit 0 clause
							        if (result == null)
							            result = new int[] {0};
        							break outer;
        						}
        						done = true;
        						break;
        					}
        				}
        				int[] results = pstatement.executeBatch();

        				if (this.executionFactory.returnsSingleUpdateCount()) {
        					// old implementation
        					for (int i = 0; i < results.length; i++) {
        						updateCount += results[i];
					        }
        					result = new int[] {updateCount};
        				} else {
        					if (result == null) {
        						result = results;
        					} else {
        						int len = result.length;
        						result = Arrays.copyOf(result, len + results.length);
        						System.arraycopy(results, 0, result, len, results.length);
        					}
				        }
        			}
        		} else {
        			bind(pstatement, translatedComm.getPreparedValues(), null);
        			updateCount = pstatement.executeUpdate();
			        result = new int[] {updateCount};
        			addStatementWarnings();
        		}
        		succeeded = true;
        	}
	        if (keyColumnNames != null && command instanceof Insert) {
		        try {
			        ResultSet keys = statement.getGeneratedKeys();
			        GeneratedKeys generatedKeys = context.returnGeneratedKeys();
			        //many databases only support returning a single generated value, but we'll still attempt to gather all
			        outer: while (keys.next()) {
				        List<Object> vals = new ArrayList<Object>(keyColumnNames.length);
				        for (int i = 0; i < keyColumnNames.length; i++) {
					        Class<?> javaType = generatedKeys.getColumnTypes()[i];
					        Object value = this.executionFactory.retrieveValue(keys, i+1, javaType);
					        if (value != null && TypeFacility.getRuntimeType(value.getClass()) != javaType) {
						        //TODO we may need to let the engine to the final conversion
						        LogManager.logDetail(LogConstants.CTX_CONNECTOR, JDBCPlugin.Util.gs(JDBCPlugin.Event.TEIID11023, javaType, keyColumnNames[i], value.getClass()));
						        continue outer;
					        }
					        vals.add(value);
				        }
				        generatedKeys.addKey(vals);
			        }
		        } catch (SQLException e) {
			        context.addWarning(e);
			        LogManager.logDetail(LogConstants.CTX_CONNECTOR, e, "Exception determining generated keys, no keys will be returned"); //$NON-NLS-1$
		        }
        	}
        } catch (SQLException err) {
        	// EXTENSION
        	// ISSUE: DVCORE-1296
        	// FIXED BY: Salvatore Raunich
        	// DATE: 2014.06.02
        	//        	throw new JDBCExecutionException(JDBCPlugin.Event.TEIID11013, err, translatedComm);
        	exc = new JDBCExecutionException(JDBCPlugin.Event.TEIID11013, err, translatedComm);
        	throw exc;
        	// END EXTENSION
        } finally {
        	if (commitType) {
        		// EXTENSION
        		// ISSUE: DVCORE-1296
        		// FIXED BY: Salvatore Raunich
        		// DATE: 2014.06.02
        		//                restoreAutoCommit(!succeeded, translatedComm);
        		restoreAutoCommit(!succeeded, translatedComm, exc);
        		// END EXTENSION
        	}
        }
    }

    /**
     * @param command
     * @return
     * @throws TranslatorException
     */
	// EXTENSION
	// ISSUE: DVCORE-1115
	// FIXED BY: Salvatore Raunich
	// DATE: 2014.09.09
//    private boolean getAutoCommit(TranslatedCommand tCommand) throws TranslatorException {
    protected boolean getAutoCommit(TranslatedCommand tCommand) throws TranslatorException {
    	// END EXTENSION
    	try {
            return connection.getAutoCommit();
        } catch (SQLException err) {
        	 throw new JDBCExecutionException(JDBCPlugin.Event.TEIID11014, err, tCommand);
        }
    }

    /**
     * If the auto comm
     * 
     * @param exceptionOccurred
     * @param command
     * @throws TranslatorException
     */
	// EXTENSION
	// ISSUE: DVCORE-1296, DVCORE-1115
	// FIXED BY: Salvatore Raunich
	// DATE: 2014.06.02, 2014.09.09
//    private void restoreAutoCommit(boolean exceptionOccurred,
//                                   TranslatedCommand tCommand) throws TranslatorException {
//    private void restoreAutoCommit(boolean exceptionOccurred,
//    		TranslatedCommand tCommand, JDBCExecutionException exc) throws TranslatorException {
    protected void restoreAutoCommit(boolean exceptionOccurred,
    		TranslatedCommand tCommand, JDBCExecutionException exc) throws TranslatorException {
    	// END EXTENSION
        try {
            if (exceptionOccurred) {
                connection.rollback();
            }
        } catch (SQLException err) {
        	// EXTENSION
        	// ISSUE: DVCORE-1296
        	// FIXED BY: Salvatore Raunich
        	// DATE: 2014.06.02
//        	throw new JDBCExecutionException(JDBCPlugin.Event.TEIID11015, err, tCommand); 
        	JDBCExecutionException currentExc = new JDBCExecutionException(JDBCPlugin.Event.TEIID11015, err, tCommand); 
        	if (exc != null) {
        		currentExc.addSuppressed(exc);
        	}
        	throw currentExc;
        	// END EXTENSION
        } finally {
        	try {
        		connection.commit(); // in JbossAs setAutocommit = true does not trigger the commit.
        		connection.setAutoCommit(true);
        	} catch (SQLException err) {
        		// EXTENSION
        		// ISSUE: DVCORE-1296
        		// FIXED BY: Salvatore Raunich
        		// DATE: 2014.06.02
//        		throw new JDBCExecutionException(JDBCPlugin.Event.TEIID11016, err, tCommand); 
        		JDBCExecutionException currentExc = new JDBCExecutionException(JDBCPlugin.Event.TEIID11016, err, tCommand); 
            	if (exc != null) {
            		currentExc.addSuppressed(exc);
            	}
            	throw currentExc;
            	//END EXTENSION
            }
        }
    }
    
    @Override
    public int[] getUpdateCounts() throws DataNotAvailableException,
    		TranslatorException {
    	return result;
    }


	public Command preTranslateUpdate(Command command) {
		return null;
	}
}
