package azkaban.test.executor;

import java.io.File;
import java.io.IOException;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;

import javax.sql.DataSource;

import junit.framework.Assert;

import org.apache.commons.dbutils.DbUtils;
import org.apache.commons.dbutils.QueryRunner;
import org.apache.commons.dbutils.ResultSetHandler;
import org.junit.BeforeClass;
import org.junit.Test;


import azkaban.executor.ExecutableFlow;
import azkaban.executor.ExecutableJobInfo;
import azkaban.executor.ExecutableNode;
import azkaban.executor.ExecutionReference;
import azkaban.executor.ExecutorLoader;
import azkaban.executor.ExecutorManagerException;
import azkaban.executor.JdbcExecutorLoader;
import azkaban.executor.Status;
import azkaban.flow.Flow;

import azkaban.utils.DataSourceUtils;
import azkaban.utils.FileIOUtils.LogData;
import azkaban.utils.JSONUtils;
import azkaban.utils.Pair;
import azkaban.utils.Props;

public class JdbcExecutorLoaderTest {
	private static boolean testDBExists;
	//@TODO remove this and turn into local host.
	private static final String host = "rpark-ld.linkedin.biz";
	private static final int port = 3306;
	private static final String database = "test";
	private static final String user = "azkaban";
	private static final String password = "azkaban";
	private static final int numConnections = 10;
	
	private File flowDir = new File("unit/executions/exectest1");
	
	@BeforeClass
	public static void setupDB() {
		DataSource dataSource = DataSourceUtils.getMySQLDataSource(host, port, database, user, password, numConnections);
		testDBExists = true;
		
		Connection connection = null;
		try {
			connection = dataSource.getConnection();
		} catch (SQLException e) {
			e.printStackTrace();
			testDBExists = false;
			DbUtils.closeQuietly(connection);
			return;
		}

		CountHandler countHandler = new CountHandler();
		QueryRunner runner = new QueryRunner();
		try {
			runner.query(connection, "SELECT COUNT(1) FROM active_executing_flows", countHandler);
		} catch (SQLException e) {
			e.printStackTrace();
			testDBExists = false;
			DbUtils.closeQuietly(connection);
			return;
		}

		try {
			runner.query(connection, "SELECT COUNT(1) FROM execution_flows", countHandler);
		} catch (SQLException e) {
			e.printStackTrace();
			testDBExists = false;
			DbUtils.closeQuietly(connection);
			return;
		}

		try {
			runner.query(connection, "SELECT COUNT(1) FROM execution_jobs", countHandler);
		} catch (SQLException e) {
			e.printStackTrace();
			testDBExists = false;
			DbUtils.closeQuietly(connection);
			return;
		}
		
		try {
			runner.query(connection, "SELECT COUNT(1) FROM execution_logs", countHandler);
		} catch (SQLException e) {
			e.printStackTrace();
			testDBExists = false;
			DbUtils.closeQuietly(connection);
			return;
		}
		
		DbUtils.closeQuietly(connection);
		
		clearDB();
	}
	
	private static void clearDB() {
		if (!testDBExists) {
			return;
		}

		DataSource dataSource = DataSourceUtils.getMySQLDataSource(host, port, database, user, password, numConnections);
		Connection connection = null;
		try {
			connection = dataSource.getConnection();
		} catch (SQLException e) {
			e.printStackTrace();
			testDBExists = false;
			DbUtils.closeQuietly(connection);
			return;
		}

		QueryRunner runner = new QueryRunner();
		try {
			runner.update(connection, "DELETE FROM active_executing_flows");
			
		} catch (SQLException e) {
			e.printStackTrace();
			testDBExists = false;
			DbUtils.closeQuietly(connection);
			return;
		}

		try {
			runner.update(connection, "DELETE FROM execution_flows");
		} catch (SQLException e) {
			e.printStackTrace();
			testDBExists = false;
			DbUtils.closeQuietly(connection);
			return;
		}
		
		try {
			runner.update(connection, "DELETE FROM execution_jobs");
		} catch (SQLException e) {
			e.printStackTrace();
			testDBExists = false;
			DbUtils.closeQuietly(connection);
			return;
		}

		try {
			runner.update(connection, "DELETE FROM execution_logs");
		} catch (SQLException e) {
			e.printStackTrace();
			testDBExists = false;
			DbUtils.closeQuietly(connection);
			return;
		}
		
		DbUtils.closeQuietly(connection);
	}
	
	@Test
	public void testUploadExecutionFlows() throws Exception {
		if (!isTestSetup()) {
			return;
		}
		
		ExecutorLoader loader = createLoader();
		ExecutableFlow flow = createExecutableFlow("exec1");

		loader.uploadExecutableFlow(flow);
		
		ExecutableFlow fetchFlow = loader.fetchExecutableFlow(flow.getExecutionId());

		// Shouldn't be the same object.
		Assert.assertTrue(flow != fetchFlow);
		Assert.assertEquals(flow.getExecutionId(), fetchFlow.getExecutionId());
		Assert.assertEquals(flow.getEndTime(), fetchFlow.getEndTime());
		Assert.assertEquals(flow.getStartTime(), fetchFlow.getStartTime());
		Assert.assertEquals(flow.getSubmitTime(), fetchFlow.getSubmitTime());
		Assert.assertEquals(flow.getFlowId(), fetchFlow.getFlowId());
		Assert.assertEquals(flow.getProjectId(), fetchFlow.getProjectId());
		Assert.assertEquals(flow.getVersion(), fetchFlow.getVersion());
		Assert.assertEquals(flow.getFailureAction(), fetchFlow.getFailureAction());
		Assert.assertEquals(new HashSet<String>(flow.getEndNodes()), new HashSet<String>(fetchFlow.getEndNodes()));
	}
	
	@Test
	public void testUpdateExecutionFlows() throws Exception {
		if (!isTestSetup()) {
			return;
		}
		
		ExecutorLoader loader = createLoader();
		ExecutableFlow flow = createExecutableFlow("exec1");

		loader.uploadExecutableFlow(flow);
		
		ExecutableFlow fetchFlow2 = loader.fetchExecutableFlow(flow.getExecutionId());
		
		fetchFlow2.setEndTime(System.currentTimeMillis());
		fetchFlow2.setStatus(Status.SUCCEEDED);
		loader.updateExecutableFlow(fetchFlow2);
		ExecutableFlow fetchFlow = loader.fetchExecutableFlow(flow.getExecutionId());
		
		// Shouldn't be the same object.
		Assert.assertTrue(flow != fetchFlow);
		Assert.assertEquals(flow.getExecutionId(), fetchFlow.getExecutionId());
		Assert.assertEquals(fetchFlow2.getEndTime(), fetchFlow.getEndTime());
		Assert.assertEquals(fetchFlow2.getStatus(), fetchFlow.getStatus());
		Assert.assertEquals(flow.getStartTime(), fetchFlow.getStartTime());
		Assert.assertEquals(flow.getSubmitTime(), fetchFlow.getSubmitTime());
		Assert.assertEquals(flow.getFlowId(), fetchFlow.getFlowId());
		Assert.assertEquals(flow.getProjectId(), fetchFlow.getProjectId());
		Assert.assertEquals(flow.getVersion(), fetchFlow.getVersion());
		Assert.assertEquals(flow.getFailureAction(), fetchFlow.getFailureAction());
		Assert.assertEquals(new HashSet<String>(flow.getEndNodes()), new HashSet<String>(fetchFlow.getEndNodes()));
	}
	
	
	@Test
	public void testUploadExecutableNode() throws Exception {
		if (!isTestSetup()) {
			return;
		}
		
		ExecutorLoader loader = createLoader();
		ExecutableFlow flow = createExecutableFlow(10, "exec1");
		flow.setExecutionId(10);
		
		File jobFile = new File(flowDir, "job10.job");
		Props props = new Props(null, jobFile);
		props.put("test","test2");
		ExecutableNode oldNode = flow.getExecutableNode("job10");
		oldNode.setStartTime(System.currentTimeMillis());
		loader.uploadExecutableNode(oldNode, props);
		
		ExecutableJobInfo info = loader.fetchJobInfo(10, "job10", 0);
		Assert.assertEquals(flow.getExecutionId(), info.getExecId());
		Assert.assertEquals(flow.getProjectId(), info.getProjectId());
		Assert.assertEquals(flow.getVersion(), info.getVersion());
		Assert.assertEquals(flow.getFlowId(), info.getFlowId());
		Assert.assertEquals(oldNode.getJobId(), info.getJobId());
		Assert.assertEquals(oldNode.getStatus(), info.getStatus());
		Assert.assertEquals(oldNode.getStartTime(), info.getStartTime());
		Assert.assertEquals("endTime = " + oldNode.getEndTime() + " info endTime = " + info.getEndTime(), oldNode.getEndTime(), info.getEndTime());
	
		// Fetch props
		Props outputProps = new Props();
		outputProps.put("hello", "output");
		oldNode.setOutputProps(outputProps);
		oldNode.setEndTime(System.currentTimeMillis());
		loader.updateExecutableNode(oldNode);

		Props fInputProps = loader.fetchExecutionJobInputProps(10, "job10");
		Props fOutputProps = loader.fetchExecutionJobOutputProps(10, "job10");
		Pair<Props,Props> inOutProps = loader.fetchExecutionJobProps(10, "job10");
		
		Assert.assertEquals(fInputProps.get("test"), "test2");
		Assert.assertEquals(fOutputProps.get("hello"), "output");
		Assert.assertEquals(inOutProps.getFirst().get("test"), "test2");
		Assert.assertEquals(inOutProps.getSecond().get("hello"), "output");
		
	}
	
	@Test
	public void testActiveReference() throws Exception {
		if (!isTestSetup()) {
			return;
		}
		
		ExecutorLoader loader = createLoader();
		ExecutableFlow flow1 = createExecutableFlow("exec1");
		loader.uploadExecutableFlow(flow1);
		ExecutionReference ref1 = new ExecutionReference(flow1.getExecutionId(), "test", 1);
		loader.addActiveExecutableReference(ref1);
		
		ExecutableFlow flow2 = createExecutableFlow("exec1");
		loader.uploadExecutableFlow(flow2);
		ExecutionReference ref2 = new ExecutionReference(flow2.getExecutionId(), "test", 1);
		loader.addActiveExecutableReference(ref2);
		
		ExecutableFlow flow3 = createExecutableFlow("exec1");
		loader.uploadExecutableFlow(flow3);
		
		Map<Integer, Pair<ExecutionReference, ExecutableFlow>> activeFlows1 = loader.fetchActiveFlows();
		ExecutableFlow flow1Result = activeFlows1.get(flow1.getExecutionId()).getSecond();
		Assert.assertNotNull(flow1Result);
		Assert.assertTrue(flow1 != flow1Result);
		Assert.assertEquals(flow1.getExecutionId(), flow1Result.getExecutionId());
		Assert.assertEquals(flow1.getEndTime(), flow1Result.getEndTime());
		Assert.assertEquals(flow1.getStartTime(), flow1Result.getStartTime());
		Assert.assertEquals(flow1.getSubmitTime(), flow1Result.getSubmitTime());
		Assert.assertEquals(flow1.getFlowId(), flow1Result.getFlowId());
		Assert.assertEquals(flow1.getProjectId(), flow1Result.getProjectId());
		Assert.assertEquals(flow1.getVersion(), flow1Result.getVersion());
		Assert.assertEquals(flow1.getFailureAction(), flow1Result.getFailureAction());
		
		ExecutableFlow flow1Result2 = activeFlows1.get(flow2.getExecutionId()).getSecond();
		Assert.assertNotNull(flow1Result2);
		Assert.assertTrue(flow2 != flow1Result2);
		Assert.assertEquals(flow2.getExecutionId(), flow1Result2.getExecutionId());
		Assert.assertEquals(flow2.getEndTime(), flow1Result2.getEndTime());
		Assert.assertEquals(flow2.getStartTime(), flow1Result2.getStartTime());
		Assert.assertEquals(flow2.getSubmitTime(), flow1Result2.getSubmitTime());
		Assert.assertEquals(flow2.getFlowId(), flow1Result2.getFlowId());
		Assert.assertEquals(flow2.getProjectId(), flow1Result2.getProjectId());
		Assert.assertEquals(flow2.getVersion(), flow1Result2.getVersion());
		Assert.assertEquals(flow2.getFailureAction(), flow1Result2.getFailureAction());
		
		loader.removeActiveExecutableReference(flow2.getExecutionId());
		Map<Integer, Pair<ExecutionReference,ExecutableFlow>> activeFlows2 = loader.fetchActiveFlows();

		Assert.assertTrue(activeFlows2.containsKey(flow1.getExecutionId()));
		Assert.assertFalse(activeFlows2.containsKey(flow3.getExecutionId()));
		Assert.assertFalse(activeFlows2.containsKey(flow2.getExecutionId()));
	}
	
	@Test
	public void testSmallUploadLog() throws ExecutorManagerException {
		File logDir = new File("unit/executions/logtest");
		File[] smalllog = {new File(logDir, "log1.log"), new File(logDir, "log2.log"), new File(logDir, "log3.log")};

		ExecutorLoader loader = createLoader();
		loader.uploadLogFile(1, "smallFiles", 0, smalllog);
		
		LogData data = loader.fetchLogs(1, "smallFiles", 0, 0, 50000);
		Assert.assertNotNull(data);
		Assert.assertEquals("Logs length is " + data.getLength(), data.getLength(), 53);
		
		System.out.println(data.toString());
	
		LogData data2 = loader.fetchLogs(1, "smallFiles", 0, 10, 20);
		System.out.println(data2.toString());
		Assert.assertNotNull(data2);
		Assert.assertEquals("Logs length is " + data2.getLength(), data2.getLength(), 20);

	}
	
	@Test
	public void testLargeUploadLog() throws ExecutorManagerException {
		File logDir = new File("unit/executions/logtest");
		
		// Multiple of 255 for Henry the Eigth
		File[] largelog = {new File(logDir, "largeLog1.log"), new File(logDir, "largeLog2.log"), new File(logDir, "largeLog3.log")};

		ExecutorLoader loader = createLoader();
		loader.uploadLogFile(1, "largeFiles",0, largelog);
		
		LogData logsResult = loader.fetchLogs(1, "largeFiles",0, 0, 64000);
		Assert.assertNotNull(logsResult);
		Assert.assertEquals("Logs length is " + logsResult.getLength(), logsResult.getLength(), 64000);
		
		LogData logsResult2 = loader.fetchLogs(1, "largeFiles",0, 1000, 64000);
		Assert.assertNotNull(logsResult2);
		Assert.assertEquals("Logs length is " + logsResult2.getLength(), logsResult2.getLength(), 64000);
		
		LogData logsResult3 = loader.fetchLogs(1, "largeFiles",0, 330000, 400000);
		Assert.assertNotNull(logsResult3);
		Assert.assertEquals("Logs length is " + logsResult3.getLength(), logsResult3.getLength(), 5493);
		
		LogData logsResult4 = loader.fetchLogs(1, "largeFiles",0, 340000, 400000);
		Assert.assertNull(logsResult4);
		
		LogData logsResult5 = loader.fetchLogs(1, "largeFiles",0, 153600, 204800);
		Assert.assertNotNull(logsResult5);
		Assert.assertEquals("Logs length is " + logsResult5.getLength(), logsResult5.getLength(), 181893);
		
		LogData logsResult6 = loader.fetchLogs(1, "largeFiles",0, 150000, 250000);
		Assert.assertNotNull(logsResult6);
		Assert.assertEquals("Logs length is " + logsResult6.getLength(), logsResult6.getLength(), 185493);
	}
	
	private ExecutableFlow createExecutableFlow(int executionId, String flowName) throws IOException {
		File jsonFlowFile = new File(flowDir, flowName + ".flow");
		@SuppressWarnings("unchecked")
		HashMap<String, Object> flowObj = (HashMap<String, Object>) JSONUtils.parseJSONFromFile(jsonFlowFile);
		
		Flow flow = Flow.flowFromObject(flowObj);
		ExecutableFlow execFlow = new ExecutableFlow(executionId, flow);

		return execFlow;
	}
	
	private ExecutableFlow createExecutableFlow(String flowName) throws IOException {
		File jsonFlowFile = new File(flowDir, flowName + ".flow");
		@SuppressWarnings("unchecked")
		HashMap<String, Object> flowObj = (HashMap<String, Object>) JSONUtils.parseJSONFromFile(jsonFlowFile);
		
		Flow flow = Flow.flowFromObject(flowObj);
		ExecutableFlow execFlow = new ExecutableFlow(flow);

		return execFlow;
	}
	
	private ExecutorLoader createLoader() {
		Props props = new Props();
		props.put("database.type", "mysql");
		
		props.put("mysql.host", host);		
		props.put("mysql.port", port);
		props.put("mysql.user", user);
		props.put("mysql.database", database);
		props.put("mysql.password", password);
		props.put("mysql.numconnections", numConnections);
		
		return new JdbcExecutorLoader(props);
	}
	
	private boolean isTestSetup() {
		if (!testDBExists) {
			System.err.println("Skipping DB test because Db not setup.");
			return false;
		}
		
		System.out.println("Running DB test because Db setup.");
		return true;
	}
	
	public static class CountHandler implements ResultSetHandler<Integer> {
		@Override
		public Integer handle(ResultSet rs) throws SQLException {
			int val = 0;
			while (rs.next()) {
				val++;
			}
			
			return val;
		}
	}
}