package meerkat.bulletinboard;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.management.ManagementFactory;
import java.lang.management.ThreadMXBean;
import java.lang.reflect.InvocationTargetException;
import java.math.BigInteger;
import java.security.InvalidKeyException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.SignatureException;
import java.security.UnrecoverableKeyException;
import java.security.cert.CertificateException;
import java.util.*;

import com.google.protobuf.BoolValue;
import com.google.protobuf.ByteString;

import com.google.protobuf.Int64Value;
import com.google.protobuf.Timestamp;
import meerkat.comm.CommunicationException;
import meerkat.comm.MessageInputStream;
import meerkat.comm.MessageOutputStream;
import meerkat.comm.MessageInputStream.MessageInputStreamFactory;
import meerkat.crypto.concrete.ECDSASignature;
import meerkat.crypto.concrete.SHA256Digest;
import meerkat.protobuf.BulletinBoardApi.*;
import meerkat.util.BulletinBoardMessageComparator;
import meerkat.util.BulletinBoardMessageGenerator;
import meerkat.util.BulletinBoardUtils;

import static org.junit.Assert.*;
import static org.hamcrest.CoreMatchers.*;
import static org.hamcrest.MatcherAssert.assertThat;

public class GenericBulletinBoardServerTest {

	protected BulletinBoardServer bulletinBoardServer;
	private GenericBulletinBoardSignature[] signers;
	private ByteString[] signerIDs;
	
	private Random random;

	private static String KEYFILE_EXAMPLE = "/certs/enduser-certs/user1-key-with-password-secret.p12";
	private static String KEYFILE_EXAMPLE3 = "/certs/enduser-certs/user3-key-with-password-shh.p12";
	
	private static String KEYFILE_PASSWORD1 = "secret";
	private static String KEYFILE_PASSWORD3 = "shh";
	
	public static String CERT1_PEM_EXAMPLE = "/certs/enduser-certs/user1.crt";
	public static String CERT3_PEM_EXAMPLE = "/certs/enduser-certs/user3.crt";

	private final int TAG_NUM = 5;		// Number of tags.
	private final int MESSAGE_NUM = 32;	// Number of messages (2^TAG_NUM).

	private String[] tags;
	private byte[][] data;

	private List<BulletinBoardMessage> batches;

	private final ThreadMXBean threadBean = ManagementFactory.getThreadMXBean(); // Used to time the tests

	private BulletinBoardMessageGenerator bulletinBoardMessageGenerator;

	private BulletinBoardDigest digest;

	private BulletinBoardMessageComparator comparator;

	/**
	 * @param bulletinBoardServer is an initialized server.
	 */
	public void init(BulletinBoardServer bulletinBoardServer) {

		System.err.println("Starting to initialize GenericBulletinBoardServerTest");
		long start = threadBean.getCurrentThreadCpuTime();

		this.bulletinBoardServer = bulletinBoardServer;
		
		signers = new GenericBulletinBoardSignature[2];
		signerIDs = new ByteString[signers.length];
        signers[0] = new GenericBulletinBoardSignature(new ECDSASignature());
        signers[1] = new GenericBulletinBoardSignature(new ECDSASignature());
		
		InputStream keyStream = getClass().getResourceAsStream(KEYFILE_EXAMPLE);
        char[] password = KEYFILE_PASSWORD1.toCharArray();

		KeyStore.Builder keyStoreBuilder = null;
		try {
			keyStoreBuilder = signers[0].getPKCS12KeyStoreBuilder(keyStream, password);

			signers[0].loadSigningCertificate(keyStoreBuilder);

			signers[0].loadVerificationCertificates(getClass().getResourceAsStream(CERT1_PEM_EXAMPLE));

			keyStream = getClass().getResourceAsStream(KEYFILE_EXAMPLE3);
			password = KEYFILE_PASSWORD3.toCharArray();

			keyStoreBuilder = signers[1].getPKCS12KeyStoreBuilder(keyStream, password);
			signers[1].loadSigningCertificate(keyStoreBuilder);

			signers[1].loadVerificationCertificates(getClass().getResourceAsStream(CERT3_PEM_EXAMPLE));

			for (int i = 0 ; i < signers.length ; i++) {
				signerIDs[i] = signers[i].getSignerID();
			}

		} catch (IOException e) {
			System.err.println("Failed reading from signature file " + e.getMessage());
			fail("Failed reading from signature file " + e.getMessage());
		} catch (CertificateException e) {
			System.err.println("Failed reading certificate " + e.getMessage());
			fail("Failed reading certificate " + e.getMessage());
		} catch (KeyStoreException e) {
			System.err.println("Failed reading keystore " + e.getMessage());
			fail("Failed reading keystore " + e.getMessage());
		} catch (NoSuchAlgorithmException e) {
			System.err.println("Couldn't find signing algorithm " + e.getMessage());
			fail("Couldn't find signing algorithm " + e.getMessage());
		} catch (UnrecoverableKeyException e) {
			System.err.println("Couldn't find signing key " + e.getMessage());
			fail("Couldn't find signing key " + e.getMessage());
		}

		// We use insecure randomness in tests for repeatability
        random = new Random(0);
		bulletinBoardMessageGenerator = new BulletinBoardMessageGenerator(random);

		digest = new GenericBulletinBoardDigest(new SHA256Digest());

		comparator = new BulletinBoardMessageComparator();

		long end = threadBean.getCurrentThreadCpuTime();
		System.err.println("Finished initializing GenericBulletinBoardServerTest");
		System.err.println("Time of operation: " + (end - start));

		// Initialize Batch variables

		batches = new ArrayList<>(10);

	}
	
	private byte randomByte(){
		return (byte) random.nextInt();
	}

	private String randomString(){
		return new BigInteger(130, random).toString(32);
	}

	/**
	 * Tests writing of several messages with multiple tags and signatures.
	 * @throws CommunicationException
	 * @throws SignatureException
	 * @throws InvalidKeyException
	 * @throws CertificateException
	 * @throws IOException
     */
	public void testInsert() {

		System.err.println("Starting to insert messages to DB");
		long start = threadBean.getCurrentThreadCpuTime();

		final int BYTES_PER_MESSAGE_DATA = 50;	// Message size.

		tags = new String[TAG_NUM];
		data = new byte[MESSAGE_NUM][BYTES_PER_MESSAGE_DATA];

		UnsignedBulletinBoardMessage.Builder unsignedMsgBuilder;
		BulletinBoardMessage.Builder msgBuilder;

		int i, j;

		// Generate random data.

		for (i = 1; i <= MESSAGE_NUM; i++) {
			for (j = 0; j < BYTES_PER_MESSAGE_DATA; j++) {
				data[i - 1][j] = randomByte();
			}
		}

		for (i = 0; i < TAG_NUM; i++) {
			tags[i] = randomString();
		}

		// Build messages.

		for (i = 1; i <= MESSAGE_NUM; i++) {
			unsignedMsgBuilder = UnsignedBulletinBoardMessage.newBuilder()
					.setData(ByteString.copyFrom(data[i - 1]))
					.setTimestamp(Timestamp.newBuilder()
							.setSeconds(i)
							.setNanos(i)
							.build());

			// Add tags based on bit-representation of message number.

			int copyI = i;
			for (j = 0; j < TAG_NUM; j++) {
				if (copyI % 2 == 1) {
					unsignedMsgBuilder.addTag(tags[j]);
				}

				copyI >>>= 1;
			}

			// Build message.

			msgBuilder = BulletinBoardMessage.newBuilder()
					.setMsg(unsignedMsgBuilder.build());

			// Add signatures.

			try {

				if (i % 2 == 1) {
					signers[0].updateContent(msgBuilder.getMsg());
					msgBuilder.addSig(signers[0].sign());

					if (i % 4 == 1) {
						signers[1].updateContent(msgBuilder.getMsg());
						msgBuilder.addSig(signers[1].sign());
					}
				}

			} catch (SignatureException e) {
				fail(e.getMessage());
			}

			// Post message.

			try {
				bulletinBoardServer.postMessage(msgBuilder.build());
			} catch (CommunicationException e) {
				fail(e.getMessage());
			}
		}

		long end = threadBean.getCurrentThreadCpuTime();
		System.err.println("Finished inserting messages to DB");
		System.err.println("Time of operation: " + (end - start));

	}

	/**
	 * Tests retrieval of messages written in {@Link #testInsert()}
	 * Only queries using one tag filter
	 */
	public void testSimpleTagAndSignature(){

		System.err.println("Starting to test tag and signature mechanism");
		long start = threadBean.getCurrentThreadCpuTime();

		List<BulletinBoardMessage> messages = new LinkedList<>();

		// Check tag mechanism

		for (int i = 0 ; i < TAG_NUM ; i++){
			
			// Retrieve messages having tag i

			try {

				MessageFilterList filterList = MessageFilterList.newBuilder()
						.addFilter(MessageFilter.newBuilder()
								.setType(FilterType.TAG)
								.setTag(tags[i])
								.build()
						)
						.build();

				ByteArrayOutputStream outputStream = new ByteArrayOutputStream();

				bulletinBoardServer.readMessages(filterList, new MessageOutputStream<BulletinBoardMessage>(outputStream));

				MessageInputStream<BulletinBoardMessage> inputStream =
						MessageInputStreamFactory.createMessageInputStream(new ByteArrayInputStream(
								outputStream.toByteArray()),
								BulletinBoardMessage.class);

				messages = inputStream.asList();

			} catch (CommunicationException | IOException e) {
				fail(e.getMessage());
				return;
			} catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
				fail(e.getMessage());
				return;
			}

			// Assert that the number of retrieved messages is correct.
			
			assertThat(messages.size(), is(MESSAGE_NUM / 2));
			
			// Assert the identity of the messages.

			for (BulletinBoardMessage msg : messages){
				
				// Assert serial number and raw data.
				
				assertThat((msg.getEntryNum() >>> i) % 2 , is((long) 1));
				assertThat(msg.getMsg().getData().toByteArray(), is(data[(int) msg.getEntryNum() - 1]));
				
				// Assert signatures.

				try {

					if (msg.getEntryNum() % 2 == 1) {
						signers[0].initVerify(msg.getSig(0));
						signers[0].updateContent(msg.getMsg());
						assertTrue("Signature did not verify!", signers[0].verify());

						if (msg.getEntryNum() % 4 == 1) {
							signers[1].initVerify(msg.getSig(1));
							signers[1].updateContent(msg.getMsg());
							assertTrue("Signature did not verify!", signers[1].verify());

							assertThat(msg.getSigCount(), is(2));
						} else {
							assertThat(msg.getSigCount(), is(1));
						}
					} else {
						assertThat(msg.getSigCount(), is(0));
					}
				} catch (Exception e) {
					fail(e.getMessage());
				}
			}
			
		}

		long end = threadBean.getCurrentThreadCpuTime();
		System.err.println("Finished testing tag and signature mechanism");
		System.err.println("Time of operation: " + (end - start));

	}

	/**
	 * Tests retrieval of messages written in {@Link #testInsert()} using multiple tags/signature filters.
	 */
	public void testEnhancedTagsAndSignatures(){

		System.err.println("Starting to test multiple tags and signatures");
		long start = threadBean.getCurrentThreadCpuTime();

		List<BulletinBoardMessage> messages;
		MessageFilterList.Builder filterListBuilder = MessageFilterList.newBuilder();

		int expectedMsgCount = MESSAGE_NUM;

		// Check multiple tag filters.

		for (int i = 0 ; i < TAG_NUM ; i++) {

			filterListBuilder.addFilter(
					MessageFilter.newBuilder()
						.setType(FilterType.TAG)
						.setTag(tags[i])
						.build()
			);

			try {

				ByteArrayOutputStream outputStream = new ByteArrayOutputStream();

				bulletinBoardServer.readMessages(filterListBuilder.build(), new MessageOutputStream<BulletinBoardMessage>(outputStream));

				MessageInputStream<BulletinBoardMessage> inputStream =
						MessageInputStreamFactory.createMessageInputStream(new ByteArrayInputStream(
								outputStream.toByteArray()),
								BulletinBoardMessage.class);

				messages = inputStream.asList();

			} catch (CommunicationException e) {
				System.err.println("Failed retrieving multi-tag messages from DB: " + e.getMessage());
				fail("Failed retrieving multi-tag messages from DB: " + e.getMessage());
				return;
			} catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException | IOException e) {
				System.err.println("Falied to read from stream while retrieving multi-tag messages: " + e.getMessage());
				fail("Falied to read from stream while retrieving multi-tag messages: " + e.getMessage());
				return;
			}

			expectedMsgCount /= 2;

			assertThat(messages.size(), is(expectedMsgCount));

			for (BulletinBoardMessage msg : messages) {
				for (int j = 0 ; j <= i ; j++) {
					assertThat((msg.getEntryNum() >>> j) % 2, is((long) 1));
				}
			}
		}

		// Check multiple signature filters.

		filterListBuilder = MessageFilterList.newBuilder()
				.addFilter(MessageFilter.newBuilder()
					.setType(FilterType.SIGNER_ID)
					.setId(signerIDs[0])
					.build())
				.addFilter(MessageFilter.newBuilder()
					.setType(FilterType.SIGNER_ID)
					.setId(signerIDs[1])
					.build());

		try {

			ByteArrayOutputStream outputStream = new ByteArrayOutputStream();

			bulletinBoardServer.readMessages(filterListBuilder.build(), new MessageOutputStream<BulletinBoardMessage>(outputStream));

			MessageInputStream<BulletinBoardMessage> inputStream =
					MessageInputStreamFactory.createMessageInputStream(new ByteArrayInputStream(
									outputStream.toByteArray()),
							BulletinBoardMessage.class);

			messages = inputStream.asList();

		} catch (CommunicationException e) {
			System.err.println("Failed retrieving multi-signature message from DB: " + e.getMessage());
			fail("Failed retrieving multi-signature message from DB: " + e.getMessage());
			return;
		} catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException | IOException e) {
			System.err.println("Falied to read from stream while retrieving multi-signature message: " + e.getMessage());
			fail("Falied to read from stream while retrieving multi-signature message: " + e.getMessage());
			return;
		}

		assertThat(messages.size(), is(MESSAGE_NUM / 4));

		for (BulletinBoardMessage message : messages) {
			assertThat(message.getEntryNum() % 4, is((long) 1));
		}

		long end = threadBean.getCurrentThreadCpuTime();
		System.err.println("Finished testing multiple tags and signatures");
		System.err.println("Time of operation: " + (end - start));

	}


	private void postAsBatch(BulletinBoardMessage message, int chunkSize, boolean close) throws CommunicationException {

		List<BatchChunk> batchChunks = BulletinBoardUtils.breakToBatch(message, chunkSize);
		BeginBatchMessage beginBatchMessage = BulletinBoardUtils.generateBeginBatchMessage(message);

		BoolValue result;

		// Begin batch

		Int64Value batchId = bulletinBoardServer.beginBatch(beginBatchMessage);

		assertThat("Was not able to open batch", batchId.getValue() != -1);

		// Post data

		BatchMessage batchMessage = BatchMessage.getDefaultInstance();

		for (int i = 0 ; i < batchChunks.size() ; i++){

			batchMessage = BatchMessage.newBuilder()
					.setBatchId(batchId.getValue())
					.setSerialNum(i)
					.setData(batchChunks.get(i))
					.build();

			result = bulletinBoardServer.postBatchMessage(batchMessage);

			assertThat("Was not able to post batch message", result.getValue(), is(true));

		}

		// Close batch
		if (close) {

			CloseBatchMessage closeBatchMessage = BulletinBoardUtils.generateCloseBatchMessage(batchId, batchChunks.size(), message);

			result = bulletinBoardServer.closeBatch(closeBatchMessage);

			assertThat("Was not able to close batch", result.getValue(), is(true));

		}

	}

	/**
	 * Posts a complete batch message
	 * @throws CommunicationException
	 */
	public void testPostBatch() throws CommunicationException, SignatureException {

		// Create data
		final int BATCH_ID = 200;
		final int DATA_SIZE = 10000;
		final int CHUNK_SIZE = 100;
		final int TAG_NUMBER = 10;

		Timestamp timestamp = Timestamp.newBuilder()
				.setSeconds(5235000)
				.setNanos(32541)
				.build();

		BulletinBoardMessage batch = bulletinBoardMessageGenerator.generateRandomMessage(signers, timestamp, DATA_SIZE, TAG_NUMBER);

		// Post batch

		postAsBatch(batch, CHUNK_SIZE, true);

		// Update locally stored batches
		batches.add(batch);

	}

	public void testReadBatch() throws CommunicationException {

		for (BulletinBoardMessage message : batches) {

			try {

				ByteArrayOutputStream outputStream = new ByteArrayOutputStream();

				digest.update(message);

				MessageID msgId = digest.digestAsMessageID();

				MessageFilterList messageFilterList = MessageFilterList.newBuilder()
						.addFilter(MessageFilter.newBuilder()
								.setType(FilterType.MSG_ID)
								.setId(msgId.getID())
								.build())
						.build();

				bulletinBoardServer.readMessages(messageFilterList, new MessageOutputStream<BulletinBoardMessage>(outputStream));

				MessageInputStream<BulletinBoardMessage> messageInputStream =
						MessageInputStreamFactory.createMessageInputStream(new ByteArrayInputStream(
										outputStream.toByteArray()),
								BulletinBoardMessage.class);

				List<BulletinBoardMessage> messageList = messageInputStream.asList();

				assertThat("No stub found for message ID " + msgId.getID().toStringUtf8(), messageList.size() == 1);

				BulletinBoardMessage stub = messageList.get(0);

				BatchQuery batchQuery =
						BatchQuery.newBuilder()
								.setMsgID(msgId)
								.setStartPosition(0)
								.build();

				bulletinBoardServer.readBatch(batchQuery, new MessageOutputStream<BatchChunk>(outputStream));

				MessageInputStream<BatchChunk> batchInputStream =
						MessageInputStreamFactory.createMessageInputStream(new ByteArrayInputStream(
								outputStream.toByteArray()),
								BatchChunk.class);

				List<BatchChunk> batchChunkList = batchInputStream.asList();

				BulletinBoardMessage retrievedMessage = BulletinBoardUtils.gatherBatch(stub, batchChunkList);

				assertThat("Non-matching batch data for batch " + msgId.getID().toStringUtf8(),
						comparator.compare(message, retrievedMessage) == 0);

			} catch (IOException | NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
				assertThat("Error reading batch data list from input stream", false);
			}

		}

	}

	public void testSyncQuery()
			throws SignatureException, CommunicationException, IOException,NoSuchMethodException, IllegalAccessException, InvocationTargetException {

		Checksum checksum = new SimpleChecksum();

		Timestamp timestamp = Timestamp.newBuilder()
				.setSeconds(1)
				.setNanos(0)
				.build();

		BulletinBoardMessage newMessage = bulletinBoardMessageGenerator.generateRandomMessage(signers, timestamp, 10, 10);

		BoolValue result = bulletinBoardServer.postMessage(newMessage);
		assertThat("Failed to post message to BB Server", result.getValue(), is(true));

		ByteArrayOutputStream outputStream = new ByteArrayOutputStream();

		digest.update(newMessage.getMsg());

		ByteString messageID = ByteString.copyFrom(digest.digest());

		MessageFilterList filterList = MessageFilterList.newBuilder()
				.addFilter(MessageFilter.newBuilder()
						.setType(FilterType.MSG_ID)
						.setId(messageID)
						.build())
				.build();

		bulletinBoardServer.readMessages(filterList, new MessageOutputStream<BulletinBoardMessage>(outputStream));

		MessageInputStream<BulletinBoardMessage> inputStream =
				MessageInputStreamFactory.createMessageInputStream(new ByteArrayInputStream(
						outputStream.toByteArray()),
						BulletinBoardMessage.class);

		long lastEntry = inputStream.asList().get(0).getEntryNum();

		SyncQuery syncQuery = SyncQuery.newBuilder()
				.setFilterList(MessageFilterList.getDefaultInstance())
				.addQuery(SingleSyncQuery.newBuilder()
						.setChecksum(2)
						.setTimeOfSync(Timestamp.newBuilder()
								.setSeconds(2)
								.setNanos(0)
								.build())
						.build())
				.build();

		SyncQueryResponse queryResponse = bulletinBoardServer.querySync(syncQuery);

		assertThat("Sync query replies with positive sync when no sync was expected", queryResponse.getLastEntryNum(), is(equalTo(-1l)));

		syncQuery = SyncQuery.newBuilder()
				.setFilterList(MessageFilterList.getDefaultInstance())
				.addQuery(SingleSyncQuery.newBuilder()
						.setChecksum(checksum.getChecksum(messageID))
						.setTimeOfSync(timestamp)
						.build())
				.build();

		queryResponse = bulletinBoardServer.querySync(syncQuery);

		assertThat("Sync query reply contained wrong last entry number", lastEntry, is(equalTo(queryResponse.getLastEntryNum())));

		assertThat("Sync query reply contained wrong timestamp", timestamp, is(equalTo(queryResponse.getLastTimeOfSync())));

	}
	
	public void close(){
		signers[0].clearSigningKey();
		signers[1].clearSigningKey();
		try {
			bulletinBoardServer.close();
		} catch (CommunicationException e) {
			System.err.println("Error closing server " + e.getMessage());
			fail("Error closing server " + e.getMessage());
		}
	}
}
