/*
 * Copyright 2018 Scytl Secure Electronic Voting SA
 *
 * All rights reserved
 *
 * See our extended copyright notice in *file 'Copyright.txt' which is part of this source code package
 */
package com.scytl.products.ov.channel.cc.commons.keymanagement;

import java.security.KeyManagementException;
import java.security.KeyPair;
import java.security.KeyStore.PasswordProtection;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.SQLIntegrityConstraintViolationException;
import java.sql.Statement;

import javax.security.auth.DestroyFailedException;
import javax.sql.DataSource;

import com.scytl.cryptolib.elgamal.bean.ElGamalPrivateKey;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;

/**
 * Implementation of {@link Database}.
 */
class DatabaseImpl implements Database {
    private final DataSource dataSource;

    private final Codec codec;

    private final Generator generator;

    private final String nodeId;

    private volatile KeyPair encryptionKeys;

    /**
     * Constructor.
     * 
     * @param dataSource
     *            the data source
     * @param codec
     *            the codec
     * @param generator
     *            the generator
     * @param nodeId
     *            the node identifier.
     */
    public DatabaseImpl(DataSource dataSource, Codec codec, Generator generator, String nodeId) {
        this.dataSource = dataSource;
        this.codec = codec;
        this.generator = generator;
        this.nodeId = nodeId;
    }

    @Override
    public ChoiceCodeKeys loadChoiceCodeKeys(String electionEventId, String verificationCardSetId)
            throws KeyNotFoundException, KeyManagementException {
        String sql = "select generation_private_key, generation_public_key, generation_signature, "
            + "decryption_private_key, decryption_public_key, decryption_signature " + "from cc_choice_code_keys "
            + "where node_id = ? and election_event_id = ? and verification_card_set_id = ?";
        try (Connection connection = dataSource.getConnection();
                PreparedStatement statement = connection.prepareStatement(sql)) {
            statement.setString(1, nodeId);
            statement.setString(2, electionEventId);
            statement.setString(3, verificationCardSetId);
            try (ResultSet resultSet = statement.executeQuery()) {
                if (!resultSet.next()) {
                    throw new KeyNotFoundException("Choice code keys not found.");
                }
                ChoiceCodeKeys.Builder builder = new ChoiceCodeKeys.Builder();
                builder.setGenerationKeys(getElGamalPrivateKey(resultSet, "generation_private_key"),
                    getElGamalPublicKey(resultSet, "generation_public_key"),
                    resultSet.getBytes("generation_signature"));
                builder.setDecryptionKeys(getElGamalPrivateKey(resultSet, "decryption_private_key"),
                    getElGamalPublicKey(resultSet, "decryption_public_key"),
                    resultSet.getBytes("decryption_signature"));
                return builder.build();
            }
        } catch (SQLException e) {
            throw new KeyManagementException("Failed to load choice code keys.", e);
        }
    }

    @Override

    public ElectionSigningKeys loadElectionSigningKeys(String electionEventId)
            throws KeyNotFoundException, KeyManagementException {
        String sql =
            "select keys, password from cc_election_signing_keys " + "where node_id = ? and election_event_id = ?";
        try (Connection connection = dataSource.getConnection();
                PreparedStatement statement = connection.prepareStatement(sql)) {
            statement.setString(1, nodeId);
            statement.setString(2, electionEventId);
            try (ResultSet resultSet = statement.executeQuery()) {
                if (!resultSet.next()) {
                    throw new KeyNotFoundException("Election signing keys not found.");
                }
                PasswordProtection password = getPassword(resultSet, "password");
                try {
                    return getElectionSigningKeys(resultSet, "keys", password);
                } finally {
                    password.destroy();
                }
            }
        } catch (SQLException | DestroyFailedException e) {
            throw new KeyManagementException("Failed to load node keys.", e);
        }
    }

    @Override
    public MixingKeys loadMixingKeys(String electionEventId, String electoralAuthorityId)
            throws KeyNotFoundException, KeyManagementException {

        String sql = "select private_key, public_key, signature " + "from cc_mixing_keys "
            + "where node_id = ? and election_event_id = ? and electoral_authority_id = ?";
        try (Connection connection = dataSource.getConnection();
                PreparedStatement statement = connection.prepareStatement(sql)) {
            statement.setString(1, nodeId);
            statement.setString(2, electionEventId);
            statement.setString(3, electoralAuthorityId);
            try (ResultSet resultSet = statement.executeQuery()) {
                if (!resultSet.next()) {
                    throw new KeyNotFoundException("Mixing keys not found.");
                }
                return new MixingKeys(getElGamalPrivateKey(resultSet, "private_key"),
                    getElGamalPublicKey(resultSet, "public_key"), resultSet.getBytes("signature"));
            }
        } catch (SQLException e) {
            throw new KeyManagementException("Failed to load mixing keys.", e);
        }
    }

    @Override
    public NodeKeys loadNodeKeys(PasswordProtection password)
            throws KeyNotFoundException, InvalidPasswordException, KeyManagementException {
        String sql = "select keys from cc_node_keys where node_id = ?";
        try (Connection connection = dataSource.getConnection();
                PreparedStatement statement = connection.prepareStatement(sql)) {
            statement.setString(1, nodeId);
            try (ResultSet resultSet = statement.executeQuery()) {
                if (!resultSet.next()) {
                    throw new KeyNotFoundException("Node keys not found.");
                }
                return getNodeKeys(resultSet, "keys", password);
            }
        } catch (SQLException e) {
            throw new KeyManagementException("Failed to load node keys.", e);
        }
    }

    @Override
    public void saveChoiceCodeKeys(String electionEventId, String verificationCardSetId, ChoiceCodeKeys keys)
            throws KeyAlreadyExistsException, KeyManagementException {
        try (Connection connection = dataSource.getConnection()) {
            if (hasChoiceCodeKeys(connection, electionEventId, verificationCardSetId)) {
                throw new KeyAlreadyExistsException("ChoiceCodeKey already exist.");
            }
            insertChoiceCodeKeys(connection, electionEventId, verificationCardSetId, keys);
        } catch (SQLException e) {
            throw new KeyManagementException("Failed to save election signing keys.", e);
        }
    }

    @Override
    public void saveElectionSigningKeys(String electionEventId, ElectionSigningKeys keys)
            throws KeyAlreadyExistsException, KeyManagementException {
        try (Connection connection = dataSource.getConnection()) {
            checkElectionSigningKeysDoNotExist(connection, electionEventId);
            PasswordProtection password = generator.generatePassword();
            try {
                insertElectionSigningKeys(connection, electionEventId, keys, password);
            } finally {
                password.destroy();
            }
        } catch (SQLException | DestroyFailedException e) {
            throw new KeyManagementException("Failed to save election signing keys.", e);
        }
    }

    @Override
    public void saveMixingKeys(String electionEventId, String electoralAuthorityId, MixingKeys keys)
            throws KeyAlreadyExistsException, KeyManagementException {
        try (Connection connection = dataSource.getConnection()) {
            if (hasMixingKeys(connection, electionEventId, electoralAuthorityId)) {
                throw new KeyAlreadyExistsException("Mixing keys already exist.");
            }
            insertMixingKeys(connection, electionEventId, electoralAuthorityId, keys);
        } catch (SQLException e) {
            throw new KeyManagementException("Failed to save mixing keys.", e);
        }
    }

    @Override
    public void saveNodeKeys(NodeKeys keys, PasswordProtection password)
            throws KeyAlreadyExistsException, KeyManagementException {
        try (Connection connection = dataSource.getConnection()) {
            lockNodeKeys(connection);
            checkNodeKeysDoNotExist(connection);
            insertNodeKeys(connection, keys, password);
        } catch (SQLException e) {
            throw new KeyManagementException("Failed to save node keys.", e);
        }
    }

    @Override
    public void setEncryptionKeys(PrivateKey privateKey, PublicKey publicKey) {
        this.encryptionKeys = new KeyPair(publicKey, privateKey);
    }

    @Override
    public boolean hasChoiceCodeKeys(String electionEventId, String verificationCardSetId)
            throws KeyManagementException {
        try (Connection connection = dataSource.getConnection()) {
            return hasChoiceCodeKeys(connection, electionEventId, verificationCardSetId);
        } catch (SQLException e) {
            throw new KeyManagementException(e);
        }
    }

    private boolean hasChoiceCodeKeys(Connection connection, String electionEventId, String verificationCardSetId)
            throws SQLException {
        String sql = "select count(*) from cc_choice_code_keys "
            + "where node_id = ? and election_event_id = ? and verification_card_set_id = ?";
        try (PreparedStatement statement = connection.prepareStatement(sql)) {
            statement.setString(1, nodeId);
            statement.setString(2, electionEventId);
            statement.setString(3, verificationCardSetId);
            try (ResultSet resultSet = statement.executeQuery()) {
                if (resultSet.next() && resultSet.getInt(1) > 0) {
                    return true;
                }
            }
        }
        return false;
    }

    private void checkElectionSigningKeysDoNotExist(Connection connection, String electionEventId)
            throws KeyAlreadyExistsException, SQLException {
        String sql = "select count(*) from cc_election_signing_keys " + "where node_id = ? and election_event_id = ?";
        try (PreparedStatement statement = connection.prepareStatement(sql)) {
            statement.setString(1, nodeId);
            statement.setString(2, electionEventId);
            try (ResultSet resultSet = statement.executeQuery()) {
                if (resultSet.next() && resultSet.getInt(1) > 0) {
                    throw new KeyAlreadyExistsException("Election signing keys already exist.");
                }
            }
        }
    }

    @Override
    public boolean hasMixingKeys(String electionEventId, String electoralAuthorityId) throws KeyManagementException {
        try (Connection connection = dataSource.getConnection()) {
            return hasMixingKeys(connection, electionEventId, electoralAuthorityId);
        } catch (SQLException e) {
            throw new KeyManagementException(e);
        }
    }

    private boolean hasMixingKeys(Connection connection, String electionEventId, String electoralAuthorityId)
            throws SQLException {
        String sql = "select count(*) from cc_mixing_keys "
            + "where node_id = ? and election_event_id = ? and electoral_authority_id = ?";
        try (PreparedStatement statement = connection.prepareStatement(sql)) {
            statement.setString(1, nodeId);
            statement.setString(2, electionEventId);
            statement.setString(3, electoralAuthorityId);
            try (ResultSet resultSet = statement.executeQuery()) {
                if (resultSet.next() && resultSet.getInt(1) > 0) {
                    return true;
                }
            }
        }
        return false;
    }

    private void checkNodeKeysDoNotExist(Connection connection) throws KeyAlreadyExistsException, SQLException {
        String sql = "select count(*) from cc_node_keys where node_id = ?";
        try (PreparedStatement statement = connection.prepareStatement(sql)) {
            statement.setString(1, nodeId);
            try (ResultSet resultSet = statement.executeQuery()) {
                if (resultSet.next() && resultSet.getInt(1) > 0) {
                    throw new KeyAlreadyExistsException("Node keys already exist.");
                }
            }
        }
    }

    private ElectionSigningKeys getElectionSigningKeys(ResultSet resultSet, String columnLabel,
            PasswordProtection password) throws KeyManagementException, SQLException {
        byte[] bytes = resultSet.getBytes(columnLabel);
        return codec.decodeElectionSigningKeys(bytes, password);
    }

    private ElGamalPrivateKey getElGamalPrivateKey(ResultSet resultSet, String columnLabel)
            throws KeyManagementException, SQLException {
        byte[] bytes = resultSet.getBytes(columnLabel);
        return codec.decodeElGamalPrivateKey(bytes, encryptionKeys.getPrivate());
    }

    private ElGamalPublicKey getElGamalPublicKey(ResultSet resultSet, String columnLabel)
            throws KeyManagementException, SQLException {
        byte[] bytes = resultSet.getBytes(columnLabel);
        return codec.decodeElGamalPublicKey(bytes);
    }

    private NodeKeys getNodeKeys(ResultSet resultSet, String columnLabel, PasswordProtection password)
            throws KeyManagementException, SQLException {
        byte[] bytes = resultSet.getBytes(columnLabel);
        return codec.decodeNodeKeys(bytes, password);
    }

    private PasswordProtection getPassword(ResultSet resultSet, String columnLabel)
            throws KeyManagementException, SQLException {
        byte[] bytes = resultSet.getBytes(columnLabel);
        return codec.decodePassword(bytes, encryptionKeys.getPrivate());
    }

    private void insertChoiceCodeKeys(Connection connection, String electionEventId, String verificationCardSetId,
            ChoiceCodeKeys keys) throws KeyAlreadyExistsException, KeyManagementException, SQLException {
        String sql = "insert into cc_choice_code_keys (" + "node_id, election_event_id, verification_card_set_id, "
            + "generation_private_key, generation_public_key, generation_signature, "
            + "decryption_private_key, decryption_public_key, decryption_signature) "
            + "values (?, ?, ?, ?, ?, ?, ?, ?, ?)";
        try (PreparedStatement statement = connection.prepareStatement(sql)) {
            statement.setString(1, nodeId);
            statement.setString(2, electionEventId);
            statement.setString(3, verificationCardSetId);

            setElGamalPrivateKey(statement, 4, keys.generationPrivateKey());
            setElGamalPublicKey(statement, 5, keys.generationPublicKey());
            statement.setBytes(6, keys.generationPublicKeySignature());

            setElGamalPrivateKey(statement, 7, keys.decryptionPrivateKey());
            setElGamalPublicKey(statement, 8, keys.decryptionPublicKey());
            statement.setBytes(9, keys.decryptionPublicKeySignature());
            statement.executeUpdate();
        } catch (SQLIntegrityConstraintViolationException e) {
            throw new KeyAlreadyExistsException("Choice code keys already exist.", e);
        }
    }

    private void insertElectionSigningKeys(Connection connection, String electionEventId, ElectionSigningKeys keys,
            PasswordProtection password) throws KeyAlreadyExistsException, KeyManagementException, SQLException {
        String sql = "insert into cc_election_signing_keys (" + "node_id, election_event_id, keys, password) "
            + "values (?, ?, ?, ?)";
        try (PreparedStatement statement = connection.prepareStatement(sql)) {
            statement.setString(1, nodeId);
            statement.setString(2, electionEventId);
            setElectionSigningKeys(statement, 3, keys, password);
            setPassword(statement, 4, password);
            statement.executeUpdate();
        } catch (SQLIntegrityConstraintViolationException e) {
            throw new KeyAlreadyExistsException("Election signing keys already exist.", e);
        }
    }

    private void insertMixingKeys(Connection connection, String electionEventId, String electoralAuthorityId,
            MixingKeys keys) throws KeyAlreadyExistsException, KeyManagementException, SQLException {
        String sql = "insert into cc_mixing_keys (" + "node_id, election_event_id, electoral_authority_id, "
            + "private_key, public_key, signature) " + "values (?, ?, ?, ?, ?, ?)";
        try (PreparedStatement statement = connection.prepareStatement(sql)) {
            statement.setString(1, nodeId);
            statement.setString(2, electionEventId);
            statement.setString(3, electoralAuthorityId);
            setElGamalPrivateKey(statement, 4, keys.privateKey());
            setElGamalPublicKey(statement, 5, keys.publicKey());
            statement.setBytes(6, keys.publicKeySignature());
            statement.executeUpdate();
        } catch (SQLIntegrityConstraintViolationException e) {
            throw new KeyAlreadyExistsException("Mixing keys already exist.", e);
        }
    }

    private void insertNodeKeys(Connection connection, NodeKeys keys, PasswordProtection password)
            throws KeyManagementException, SQLException {
        String sql = "insert into cc_node_keys (node_id, keys) values (?, ?)";
        try (PreparedStatement statement = connection.prepareStatement(sql)) {
            statement.setString(1, nodeId);
            setNodeKeys(statement, 2, keys, password);
            statement.executeUpdate();
        } catch (SQLIntegrityConstraintViolationException e) {
            throw new KeyAlreadyExistsException("Node keys already exist.", e);
        }
    }

    private void lockNodeKeys(Connection connection) throws SQLException {
        String sql = "lock table cc_node_keys in exclusive mode";
        try (Statement statement = connection.createStatement()) {
            statement.executeUpdate(sql);
        }
    }

    private void setElectionSigningKeys(PreparedStatement statement, int index, ElectionSigningKeys keys,
            PasswordProtection password) throws KeyManagementException, SQLException {
        byte[] bytes = codec.encodeElectionSigningKeys(keys, password);
        statement.setBytes(index, bytes);
    }

    private void setElGamalPrivateKey(PreparedStatement statement, int index, ElGamalPrivateKey key)
            throws KeyManagementException, SQLException {
        byte[] bytes = codec.encodeElGamalPrivateKey(key, encryptionKeys.getPublic());
        statement.setBytes(index, bytes);
    }

    private void setElGamalPublicKey(PreparedStatement statement, int index, ElGamalPublicKey key)
            throws KeyManagementException, SQLException {
        byte[] bytes = codec.encodeElGamalPublicKey(key);
        statement.setBytes(index, bytes);
    }

    private void setNodeKeys(PreparedStatement statement, int index, NodeKeys keys, PasswordProtection password)
            throws KeyManagementException, SQLException {
        byte[] bytes = codec.encodeNodeKeys(keys, password);
        statement.setBytes(index, bytes);
    }

    private void setPassword(PreparedStatement statement, int index, PasswordProtection password)
            throws KeyManagementException, SQLException {
        byte[] bytes = codec.encodePassword(password, encryptionKeys.getPublic());
        statement.setBytes(index, bytes);
    }
}
