package com.robonobo.wang.client;



import java.io.ByteArrayOutputStream;

import java.io.File;

import java.io.FileInputStream;

import java.io.FileOutputStream;

import java.io.IOException;

import java.security.GeneralSecurityException;

import java.security.MessageDigest;

import java.security.NoSuchAlgorithmException;

import java.util.HashMap;

import java.util.LinkedList;

import java.util.List;

import java.util.Map;



import javax.crypto.Cipher;

import javax.crypto.spec.SecretKeySpec;



import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;



import com.robonobo.wang.proto.WangProtocol.CoinMsg;



/**

 * Stores coins securely on the filesystem. If there are no public methods

 * currently being executed, the file store is guaranteed to be in a recoverable

 * state.

 */

public class CoinStore {

	private Log log = LogFactory.getLog(getClass());

	private File storageDir;

	private Cipher encryptCipher;

	private Cipher decryptCipher;

	private Map<Integer, List<String>> coinIds = new HashMap<Integer, List<String>>();



	public CoinStore(File storageDir, String password) {

		// Create our encryption ciphers

		try {

			byte[] keyArr = generateKey(password, 16);

			SecretKeySpec sekritKey = new SecretKeySpec(keyArr, "AES");

			encryptCipher = Cipher.getInstance("AES");

			encryptCipher.init(Cipher.ENCRYPT_MODE, sekritKey);

			decryptCipher = Cipher.getInstance("AES");

			decryptCipher.init(Cipher.DECRYPT_MODE, sekritKey);

		} catch (Exception e) {

			throw new RuntimeException(e);

		}

		// Load coins from dir

		this.storageDir = storageDir;

		if(!storageDir.exists())

			storageDir.mkdirs();

		log.info("CoinStore loading persisted coins from "+storageDir.getAbsolutePath());

		double coinVal = 0;

		int numCoins = 0;

		for (File coinFile : storageDir.listFiles()) {

			try {

				CoinMsg coin = loadCoinFromFile(coinFile.getName());

				if(!coinIds.containsKey(coin.getDenom()))

					coinIds.put(coin.getDenom(), new LinkedList<String>());

				coinIds.get(coin.getDenom()).add(coin.getCoinId().toString());

				coinVal += getDenomValue(coin.getDenom());

				numCoins++;

			} catch (Exception e) {

				log.error("CoinStore encountered error ("+e.getClass().getName()+") loading coin. Skipping.");

			}

		}

		log.info("CoinStore loaded "+numCoins+" coins, value="+coinVal);

	}



	public synchronized int numCoins(int denom) {

		return (coinIds.containsKey(denom)) ? coinIds.get(denom).size() : 0;

	}



	public synchronized CoinMsg getCoin(int denom) throws CoinStoreException {

		if(!coinIds.containsKey(denom) || coinIds.get(denom).size() == 0)

			return null;

		String coinId = coinIds.get(denom).remove(0);

		CoinMsg coin;

		try {

			coin = loadCoinFromFile(coinId);

		} catch (Exception e) {

			throw new CoinStoreException(e);

		}

		new File(storageDir, coinId).delete();

		return coin;

	}

	

	public synchronized void putCoin(CoinMsg coin) throws CoinStoreException {

		if(!coinIds.containsKey(coin.getDenom()))

			coinIds.put(coin.getDenom(), new LinkedList<String>());

		coinIds.get(coin.getDenom()).add(coin.getCoinId().toString());

		try {

			saveCoinToFile(coin);

		} catch (Exception e) {

			throw new CoinStoreException(e);

		}

	}



	private void saveCoinToFile(CoinMsg coin) throws IOException, GeneralSecurityException {

		byte[] encArr = encryptCipher.doFinal(coin.toByteArray());

		File file = new File(storageDir, coin.getCoinId().toString());

		FileOutputStream fos = new FileOutputStream(file);

		fos.write(encArr);

		fos.close();

	}





	private CoinMsg loadCoinFromFile(String coinId) throws IOException, GeneralSecurityException {

		File file = new File(storageDir, coinId);

		FileInputStream fis = new FileInputStream(file);

		byte[] readArr = new byte[1024];

		ByteArrayOutputStream baos = new ByteArrayOutputStream();

		int numRead;

		while((numRead = fis.read(readArr)) > 0) {

			baos.write(readArr, 0, numRead);

		}

		fis.close();

		byte[] plainArr = decryptCipher.doFinal(baos.toByteArray());

		return CoinMsg.parseFrom(plainArr);

	}



	private byte[] generateKey(String password, int keyLength) {

		MessageDigest m;

		try {

			m = MessageDigest.getInstance("sha-256");

		} catch (NoSuchAlgorithmException e) {

			throw new RuntimeException();

		}

		m.update(password.getBytes());

		byte[] key = new byte[keyLength];

		System.arraycopy(m.digest(), 0, key, 0, keyLength);

		return key;

	}



	private double getDenomValue(Integer denom) {

		return Math.pow(2, denom);

	}

}