package aalg.crypto;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import aalg.algebra.Utils;

/**
 * Controller of anonymous authentication protocol.
 * @author Ing. Petr Ležák
 *
 * @param <PUB> type of public key
 * @param <PAYLOAD> type of payload
 */
public class AnonymousController<PUB extends IPublicKey, PAYLOAD> extends AnonymousBase {
	
	private static final String INFO = "Info";
	
	private final IRandomGenerator rng;
	private final IAssymetricCipher<? extends IPrivateKey, PUB> assymetricCipher;
	private int responseLength;
	
	private long nextPayloadId;
	private long nextChallengeId;
	
	private static class PayloadRecord<PUB extends IPublicKey, PAYLOAD> {
		private final AnonymousVerificationValues<PUB> verificationValues;
		private final PAYLOAD payload;
		
		public PayloadRecord(final AnonymousVerificationValues<PUB> verificationValues, final PAYLOAD payload) {
			this.verificationValues = verificationValues;
			this.payload = payload;
		}
	}
	
	private static class ChallengeRecord {
		private final long payloadId;
		private final byte[] response;
		
		public ChallengeRecord(final long payloadId, final byte[] response) {
			this.payloadId = payloadId;
			this.response = response;
		}
	}
	
	private final Map<Long, PayloadRecord<PUB, PAYLOAD>> payloadMap;
	private final Map<Long, ChallengeRecord> challengeMap;
	
	public AnonymousController(final IRandomGenerator rng, final IAssymetricCipher<? extends IPrivateKey, PUB> assymetricCipher) {
		this.rng = rng;
		this.assymetricCipher = assymetricCipher;
		
		this.payloadMap = new HashMap<Long, PayloadRecord<PUB, PAYLOAD>>();
		this.challengeMap = new HashMap<Long, ChallengeRecord>();
		
		this.nextPayloadId = 0;
		this.nextChallengeId = 0;
		this.responseLength = 32;
	}
	
	public long addPayload (final PAYLOAD payload, final AnonymousVerificationValues<PUB> verificationValues) {
		final long payloadId = nextPayloadId;
		nextPayloadId++;
		
		payloadMap.put(payloadId, new PayloadRecord<PUB, PAYLOAD>(verificationValues, payload));
		
		return payloadId;
	}
	
	public Set<Long> getPayloadIds() {
		return new HashSet<Long>(payloadMap.keySet());
	}
	
	public AnonymousChallenge getChallenge(final long payloadId) throws IOException {
		final byte[] response = new byte[responseLength];
		rng.fillArray(response, 0, responseLength);
		
		final PayloadRecord<PUB, PAYLOAD> payloadRecord = payloadMap.get(payloadId);
		
		if (payloadRecord != null) {
			final PUB challengeKey = payloadRecord.verificationValues.getTmpPubKey();
			final byte[] challenge = createChallenge (response, INFO, challengeKey);
			
			final long challengeId = nextChallengeId;
			nextChallengeId++;
			
			challengeMap.put(challengeId, new ChallengeRecord(payloadId, response));
			
			return new AnonymousChallenge(challengeId, payloadRecord.verificationValues.getEncryptedMessage(), challenge);
		}
		
		throw new RuntimeException("Cannot get challenge");
	}
	
	private byte[] createChallenge(final byte[] response, final String info, final PUB challengeKey) throws IOException {
		final ByteArrayOutputStream targerStream = new ByteArrayOutputStream();
		assymetricCipher.setPublicKey(challengeKey);
		
		final IStreamEncryptor encryptor = assymetricCipher.getEncryptor();
		encryptor.initialize(targerStream);
		
		final OutputStream sourceStream = encryptor.getSourceStream();
		
		Utils.writeByteArrayWithSize(sourceStream, response, RESPONSE_LENGTH_BYTES);
		
		final byte[] infoBytes = info.getBytes(INFO_CHARSET);
		Utils.writeByteArrayWithSize(sourceStream, infoBytes, INFO_LENGTH_BYTES);		
		
		encryptor.finish();
		
		return targerStream.toByteArray();
	}

	
	public PAYLOAD getPayload(final long challengeId, final byte[] response) {
		final ChallengeRecord challenge = challengeMap.get(challengeId);
		
		if (challenge != null && Arrays.equals(response, challenge.response)) {
			challengeMap.remove(challengeId);
			
			final PayloadRecord<PUB, PAYLOAD> payloadRecord = payloadMap.remove(challenge.payloadId);
			
			if (payloadRecord != null)
				return payloadRecord.payload;
		}
		
		throw new RuntimeException("Error in authentication");
	}

}
