package io.irontest.core.teststep;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.module.SimpleModule;
import io.irontest.db.RetainingColumnOrderResultSetMapper;
import io.irontest.db.SQLStatementType;
import io.irontest.models.OracleTIMESTAMPTZSerializer;
import io.irontest.models.endpoint.Endpoint;
import io.irontest.models.teststep.Teststep;
import io.irontest.utils.IronTestUtils;
import org.jdbi.v3.core.Handle;
import org.jdbi.v3.core.Jdbi;
import org.jdbi.v3.core.statement.Query;
import org.jdbi.v3.core.statement.Script;
import org.jdbi.v3.core.statement.StatementContext;
import org.jdbi.v3.core.statement.StatementCustomizer;

import java.sql.PreparedStatement;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class DBTeststepRunner extends TeststepRunner {
    private static ObjectMapper jacksonObjectMapper = new ObjectMapper();

    static {
        jacksonObjectMapper.setDateFormat(new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss"));

        //  settings for Oracle
        //  for ResultSet.getObject to return java.sql.Timestamp instead of oracle.sql.TIMESTAMP which Jackson does not know how to serialize.
        System.getProperties().setProperty("oracle.jdbc.J2EE13Compliant", "true");
        //  register custom JSON serializer for Oracle TIMESTAMPTZ class
        try {
            Class clazz = Class.forName("oracle.sql.TIMESTAMPTZ");
            SimpleModule module = new SimpleModule("OracleModule");
            module.addSerializer(clazz, new OracleTIMESTAMPTZSerializer(clazz));
            jacksonObjectMapper.registerModule(module);
        } catch (ClassNotFoundException e) {
            //  do nothing if the TIMESTAMPTZ class does not exist
        }
    }

    public BasicTeststepRun run() throws Exception {
        Teststep teststep = getTeststep();
        BasicTeststepRun basicTeststepRun = new BasicTeststepRun();
        DBAPIResponse response = new DBAPIResponse();
        String request = (String) teststep.getRequest();

        List<String> statements = IronTestUtils.getStatements(request);
        sanityCheckTheStatements(statements);

        Endpoint endpoint = teststep.getEndpoint();
        Jdbi jdbi;
        if (endpoint.getUsername() == null) {
            jdbi = Jdbi.create(endpoint.getUrl());
        } else {
            jdbi = Jdbi.create(endpoint.getUrl(), endpoint.getUsername(), getDecryptedEndpointPassword());
        }
        Handle handle = jdbi.open();
        if (SQLStatementType.isSelectStatement(statements.get(0))) {    //  the request is a select statement
            RetainingColumnOrderResultSetMapper resultSetMapper = new RetainingColumnOrderResultSetMapper();
            //  use statements.get(0) instead of the raw request, as Oracle does not support trailing semicolon in select statement
            Query query = handle.createQuery(statements.get(0))
                    .setMaxRows(5000);          //  limit the number of returned rows
            //  obtain columnNames in case the query returns no row
            final List<String> columnNames = new ArrayList<String>();
            query.addCustomizer(new StatementCustomizer() {
                public void afterExecution(PreparedStatement stmt, StatementContext ctx) throws SQLException {
                    ResultSetMetaData metaData = stmt.getMetaData();
                    for (int i = 1; i <= metaData.getColumnCount(); i++) {
                        columnNames.add(metaData.getColumnLabel(i));
                    }
                }
            });

            List<Map<String, Object>> rows = query.map(resultSetMapper).list();
            response.setColumnNames(columnNames);
            response.setRowsJSON(jacksonObjectMapper.writeValueAsString(rows));
        } else {                                          //  the request is one or more non-select statements
            Script script = handle.createScript(request);
            int[] returnValues = script.execute();
            StringBuilder sb = new StringBuilder();
            for (int i = 0; i < returnValues.length; i++) {
                String statementType = SQLStatementType.getByStatement(statements.get(i)).toString();
                sb.append(returnValues[i]).append(" row(s) ").append(statementType.toLowerCase())
                    .append(statementType.endsWith("E") ? "d" : "ed").append("\n");
                response.setStatementExecutionResults(sb.toString());
            }
        }

        handle.close();

        basicTeststepRun.setResponse(response);
        return basicTeststepRun;
    }

    /**
     * Throw exception if the statements are not supported.
     * @param statements
     * @throws Exception
     */
    private void sanityCheckTheStatements(List<String> statements) throws Exception {
        if (statements.size() == 0) {
            throw new Exception("No SQL statement to run.");
        }

        int selectStatementCount = 0;
        boolean nonSelectStatementExists = false;
        for (String statement: statements) {
            if (!(SQLStatementType.isSelectStatement(statement) ||
                    SQLStatementType.isInsertStatement(statement) ||
                    SQLStatementType.isUpdateStatement(statement) ||
                    SQLStatementType.isDeleteStatement(statement))) {
                throw new Exception("Only " + SQLStatementType.SELECT + ", " + SQLStatementType.INSERT + ", " +
                        SQLStatementType.UPDATE + " and " + SQLStatementType.DELETE + " statements are supported.");
            }
            if (SQLStatementType.isSelectStatement(statement)) {
                selectStatementCount++;
            } else {
                nonSelectStatementExists = true;
            }
        }
        if (selectStatementCount > 1) {
            throw new Exception("At most one " + SQLStatementType.SELECT + " statement is supported.");
        }
        if (selectStatementCount == 1 && nonSelectStatementExists) {
            throw new Exception("Mixture of " + SQLStatementType.SELECT + " and non-" + SQLStatementType.SELECT + " statements are not supported.");
        }
    }
}