package dslab.client;

import java.io.*;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;
import java.util.Base64;
import java.util.logging.Logger;

import dslab.ComponentFactory;
import dslab.exception.FailedVerificationException;
import dslab.util.Config;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class MessageClient implements IMessageClient, Runnable {
    private static final Logger logger = Logger.getLogger(MessageClient.class.getName());

    private final String transferHost;
    private final int transferPort;
    private final String mailboxHost;
    private final int mailboxPort;

    private final String transferEmail;

    private final String mailboxUser;
    private final String mailboxPassword;

    private Socket dmtpSocket;
    private Socket dmapSocket;

    private PrintWriter dmtpOut;
    private BufferedReader dmtpIn;
    private PrintWriter dmapOut;
    private BufferedReader dmapIn;

    private Cipher aesEncryptCipher;
    private Cipher aesDecryptCipher;
    private byte[] challenge;

    /**
     * Creates a new client instance.
     *
     * @param componentId the id of the component that corresponds to the Config resource
     * @param config      the component config
     * @param in          the input stream to read console input from
     * @param out         the output stream to write console output to
     */
    public MessageClient(String componentId, Config config, InputStream in, PrintStream out) {
        this.transferHost = config.getString("transfer.host");
        this.transferPort = config.getInt("transfer.port");
        this.mailboxHost = config.getString("mailbox.host");
        this.mailboxPort = config.getInt("mailbox.port");
        this.transferEmail = config.getString("transfer.email");
        this.mailboxUser = config.getString("mailbox.user");
        this.mailboxPassword = config.getString("mailbox.password");

        logger.info(String.format("TransferHost: %s\nTransferPort: %d\nMailboxHost: %s\nMailboxPort: %d\nTransferEmail: %s\nMailboxUser: %s\nMailboxPassword: %s",
                transferHost, transferPort, mailboxHost, mailboxPort, transferEmail, mailboxUser, mailboxPassword));
    }

    @Override
    public void run() {
        try {
            logger.info("Starting connection to MailboxHost on " + this.mailboxHost + " on port " + this.mailboxPort);
            this.dmapSocket = new Socket(this.mailboxHost, this.mailboxPort);
            this.dmapIn = new BufferedReader(new InputStreamReader(this.dmapSocket.getInputStream()));
            this.dmapOut = new PrintWriter(this.dmapSocket.getOutputStream(), true);

            String input = null;
            String message = null;

            input = dmapIn.readLine();
            logger.info("Received message from server: " + input);
            if (!input.startsWith("ok DMAP2.0"))
                shutdown();

            startSecure();
            loginUser();

            while (!Thread.currentThread().isInterrupted() && ((input = this.dmapIn.readLine()) != null)) {
                // TODO implement commands
            }
        } catch (IOException e) {
            logger.severe("Could not connect to MailboxHost " + this.mailboxHost + " on port " + this.mailboxPort);
            shutdown();
        } catch (IllegalBlockSizeException | BadPaddingException e) {
            logger.severe(e.getMessage());
            shutdown();
        }
    }

    private void loginUser() throws BadPaddingException, IllegalBlockSizeException, IOException {
        String message;
        logger.info("Trying to log in using " + this.mailboxUser + " " + this.mailboxPassword);
        message = "login " + this.mailboxUser + " " + this.mailboxPassword;
        this.dmapOut.println(getAesCiphertext(message));
        String response = getAesPlaintext(this.dmapIn.readLine());
        if (!response.startsWith("ok DMAP2.0"))
            shutdown();
    }

    private void startSecure() throws IOException {
        String input;
        String componentId;
        PublicKey serverPublicKey;
        this.dmapOut.println("startsecure");
        logger.info("Sent command 'startsecure'");
        input = this.dmapIn.readLine();
        logger.info("Server's response: " + input);
        if (input.startsWith("ok") && (input.split("\\s+").length == 2)) {
            // Get the component-id from the server
            componentId = input.split("\\s+")[1];
            try {
                // Attempt to read server public key from file called <component-id>_pub.der
                byte[] keyBytes = Files.readAllBytes(Paths.get("keys", "client", componentId + "_pub.der"));
                logger.info("Read bytes from path " + Paths.get("keys", "client", componentId + "_pub.der") + ": " + Arrays.toString(keyBytes));
                // Create X509 spec object from key
                X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
                // Create generator for RSA scheme
                KeyFactory kf = KeyFactory.getInstance("RSA");
                // Make generator generate public key from X509 spec
                serverPublicKey = kf.generatePublic(spec);
                logger.info("Server's Public Key: " + serverPublicKey);
                String clientChallenge = generateChallengeMessage(serverPublicKey);
                // Send clientChallenge to server
                logger.info("Send clientchallenge to Server: " + clientChallenge);
                this.dmapOut.println(clientChallenge);
                // Receive AES encrypted message saying "ok <client-challenge>"
                String response = this.dmapIn.readLine();
                // Compare received client challenge with generated client challenge
                verifyChallenge(response);
                // Answer with AES encrypted "ok" if matching and use AES cipher for subsequent communication
                logger.info("Send ciphered 'ok' to Server: " + getAesCiphertext("ok"));
                this.dmapOut.println(getAesCiphertext("ok"));
            } catch (NoSuchAlgorithmException | InvalidKeySpecException | IllegalBlockSizeException | BadPaddingException | FailedVerificationException e) {
                logger.severe(e.getMessage());
                shutdown();
            }
        } else {
            logger.severe("MailboxHost " + mailboxHost + " on port " + mailboxPort + " did not respond with 'ok' <component-id>");
            shutdown();
        }
    }

    /**
     * Takes a plaintext message, encrypts and encodes it and returns the encoded ciphertext ready for sending.
     *
     * @param message Plaintext message
     * @return B64 encoded and AES encrypted ciphertext
     */
    private String getAesCiphertext(String message) throws IllegalBlockSizeException, BadPaddingException {
        byte[] cipherText = aesEncryptCipher.doFinal(message.getBytes(StandardCharsets.UTF_8));
        return Base64.getEncoder().encodeToString(cipherText);
    }

    /**
     * Takes an AES encrypted message, decrypts and decodes it and returns the plaintext.
     *
     * @param message B64 encoded and AES encrypted message
     * @return B64 decoded and AES decrypted plaintext
     */
    private String getAesPlaintext(String message) throws BadPaddingException, IllegalBlockSizeException {
        byte[] cipherText = Base64.getDecoder().decode(message.getBytes(StandardCharsets.UTF_8));
        return new String(aesDecryptCipher.doFinal(cipherText));
    }

    private void verifyChallenge(String response) throws FailedVerificationException {
        // Decrypt to base64 encoded byte array
        String plainText;
        try {
            plainText = new String(aesDecryptCipher.doFinal(Base64.getDecoder().decode(response)));
            logger.info("Decrypted AES challenge: " + plainText);
        } catch (IllegalBlockSizeException | BadPaddingException e) {
            logger.severe("Error during decryption of client challenge. Aborting...");
            shutdown();
            return;
        }
        // Decode <client-challenge>
        if (plainText.split("\\s+").length != 2)
            shutdown();
        byte[] clientChallenge = Base64.getDecoder().decode(plainText.split("\\s+")[1]);
        // Compare received challenge with sent

        if (!Arrays.equals(clientChallenge, this.challenge))
            throw new FailedVerificationException("Server's clientChallenge " + new String(clientChallenge) +
                    " does not match sent clientChallenge " + new String(this.challenge));
    }

    private SecretKeySpec generateSecretKey() {
        try {
            KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");
            keyGenerator.init(256, SecureRandom.getInstanceStrong());
            return new SecretKeySpec(keyGenerator.generateKey().getEncoded(), "AES");
        } catch (NoSuchAlgorithmException e) {
            logger.severe("This should not be happening!");
            e.printStackTrace();
            shutdown();
        }
        return null;
    }

    private IvParameterSpec generateIv() {
        // Size of IV corresponds to AES block size (=128bits=16bytes)
        byte[] iv = new byte[16];
        new SecureRandom().nextBytes(iv);
        return new IvParameterSpec(iv);
    }

    private void setAesCiphers(SecretKeySpec secretKey, IvParameterSpec iv) {
        try {
            this.aesEncryptCipher = Cipher.getInstance("AES/CTR/NoPadding");
            this.aesDecryptCipher = Cipher.getInstance("AES/CTR/NoPadding");
        } catch (NoSuchAlgorithmException | NoSuchPaddingException e) {
            logger.severe("This should not be happening!");
            e.printStackTrace();
            shutdown();
        }

        try {
            this.aesEncryptCipher.init(Cipher.ENCRYPT_MODE, secretKey, iv);
            this.aesDecryptCipher.init(Cipher.DECRYPT_MODE, secretKey, iv);
        } catch (InvalidKeyException | InvalidAlgorithmParameterException e) {
            logger.severe("This should not be happening!");
            e.printStackTrace();
            shutdown();
        }
    }

    private byte[] generateChallenge(PublicKey serverPublicKey) {
        SecureRandom secureRandom = new SecureRandom();
        // Generate new random 32 byte challenge
        this.challenge = new byte[32];
        secureRandom.nextBytes(this.challenge);
        return this.challenge;
    }

    private String generateChallengeMessage(PublicKey serverPublicKey) {
        SecureRandom secureRandom = new SecureRandom();
        byte[] clearTextChallenge = generateChallenge(serverPublicKey);
        SecretKeySpec secretKeySpec = generateSecretKey();
        assert secretKeySpec != null;
        IvParameterSpec iv = generateIv();
        // Save AES cipher for subsequent communication
        setAesCiphers(secretKeySpec, iv);
        // Encode parameters to base64
        String clearTextChallengeEncoded = Base64.getEncoder().encodeToString(clearTextChallenge);
        String secretKeyEncoded = Base64.getEncoder().encodeToString(secretKeySpec.getEncoded());
        String ivEncoded = Base64.getEncoder().encodeToString(iv.getIV());
        // Concatenate command and parameters (challenge, secretKey and IV)
        String concatenated = "ok " + clearTextChallengeEncoded + " " + secretKeyEncoded + " " + ivEncoded;
        // Encrypt "<base64Encoded>"
        Cipher cipher;
        byte[] cipherTextChallenge;
        try {
            cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            cipher.init(Cipher.ENCRYPT_MODE, serverPublicKey);
            cipher.update((concatenated).getBytes(StandardCharsets.UTF_8));
            cipherTextChallenge = cipher.doFinal();
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException | BadPaddingException e) {
            logger.severe("This should not be happening!");
            e.printStackTrace();
            shutdown();
            return null;
        }
        // Return base64 encoded cipherMessage
        return (Base64.getEncoder().encodeToString(cipherTextChallenge));
    }

    @Override
    public void inbox() {

    }

    @Override
    public void delete(String id) {

    }

    @Override
    public void verify(String id) {

    }

    @Override
    public void msg(String to, String subject, String data) {
        try {
            logger.info("Starting connection to TransferServer on host " + this.transferHost + " on port " + this.transferPort);
            this.dmtpSocket = new Socket(this.transferHost, this.transferPort);
            this.dmtpIn = new BufferedReader(new InputStreamReader(this.dmtpSocket.getInputStream()));
            this.dmtpOut = new PrintWriter(this.dmtpSocket.getOutputStream(), true);
        } catch (IOException e) {
            logger.severe("Could not connect to TransferHost " + transferHost + " on port " + transferPort);
            shutdown();
        }
    }

    @Override
    public void shutdown() {
        if (dmtpSocket != null) {
            try {
                this.dmtpSocket.close();
                this.dmtpOut.close();
                this.dmtpIn.close();
            } catch (IOException e) {
                logger.severe("Could not close connection to TransferHost");
                e.printStackTrace();
            }
        }

        if (dmapSocket != null) {
            try {
                this.dmapSocket.close();
                this.dmapOut.close();
                this.dmapIn.close();
            } catch (IOException e) {
                logger.severe("Could not close connection to MailboxHost");
                e.printStackTrace();
            }
        }

        Thread.currentThread().interrupt();
    }

    public static void main(String[] args) throws Exception {
        IMessageClient client = ComponentFactory.createMessageClient(args[0], System.in, System.out);
        client.run();
    }
}