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.Level;
import java.util.logging.Logger;

import at.ac.tuwien.dsg.orvell.Shell;
import at.ac.tuwien.dsg.orvell.StopShellException;
import at.ac.tuwien.dsg.orvell.annotation.Command;
import dslab.ComponentFactory;
import dslab.Email;
import dslab.Message;
import dslab.exception.FailedVerificationException;
import dslab.exception.MalformedInputException;
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 final Shell shell;

    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");

        this.shell = new Shell(in, out);
        this.shell.register(this);
        this.shell.setPrompt(componentId + "> ");

        logger.fine(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;

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

            startSecure();
            loginUser();
            shell.run();

        } 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"))
            shutdown();
    }

    private void startSecure() throws IOException {
        String input;
        String componentId;
        PublicKey serverPublicKey;
        this.dmapOut.println("startsecure");
        logger.finer("Sent command 'startsecure'");
        input = this.dmapIn.readLine();
        logger.finer("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.finer("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);
                String clientChallenge = generateChallengeMessage(serverPublicKey);
                // Send clientChallenge to server
                logger.finer("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.finer("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));
    }

    /**
     * Takes a server's AES encrypted challenge and compares it against the one that was sent initially.
     *
     * @param response The server's encoded and encrypted client challenge.
     * @throws FailedVerificationException Thrown if the challenges do not match.
     */
    private void verifyChallenge(String response) throws FailedVerificationException {
        // Decrypt to base64 encoded byte array
        logger.info("Verifying challenge...");
        String plainText;
        try {
            plainText = new String(aesDecryptCipher.doFinal(Base64.getDecoder().decode(response)));
            logger.finer("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));
    }

    /**
     * Generates a new random 256 bit AES secret key.
     *
     * @return Either shuts down on error or returns the generated key.
     */
    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;
    }

    /**
     * Generates a new random 128 bit initialization vector
     *
     * @return The generated IV.
     */
    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);
    }

    /**
     * Sets the global encryption and decryption ciphers (aesEncryptCipher, aesDecryptCipher).
     *
     * @param secretKey A previously generated 256 bit AES secret key.
     * @param iv        A previously generated 128 bit AES initialization vector.
     */
    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();
        }
    }

    /**
     * Generates a new random 256 bit challenge.
     *
     * @return The generated challenge.
     */
    private byte[] generateChallenge() {
        SecureRandom secureRandom = new SecureRandom();
        // Generate new random 32 byte challenge
        this.challenge = new byte[32];
        secureRandom.nextBytes(this.challenge);
        return this.challenge;
    }

    /**
     * Generates the full challenge message to be sent to the server.
     * <p>
     * The challenge message is of the format:
     * ok <client-challenge> <secret-key> <iv>
     * The parameters are base64 encoded individually, then they are concatenated:
     * ok <base64-client-challenge> <base64-secret-key> <base64-iv>
     * The whole string is then AES encrypted and the result base64 encoded again.
     *
     * @return A base64 encoded full client challenge.
     */
    private String generateChallengeMessage(PublicKey serverPublicKey) {
        SecureRandom secureRandom = new SecureRandom();
        byte[] clearTextChallenge = generateChallenge();
        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));
    }

    @Command
    @Override
    public void inbox() {
        logger.info("Received 'inbox' command");

    }

    @Command
    @Override
    public void delete(String id) {
        logger.info("Received 'delete' command for id " + id);

    }

    /**
     * Pulls a message by id from the MailboxServer and compares the calculated hash against the one received via the
     * 'hash' field.
     *
     * @param id The message id.
     */
    @Command
    @Override
    public void verify(String id) {
        logger.info("Received 'verify' command for id " + id);
        Message msg = new Message();
        String response;
        // Get message from MailboxServer
        try {
            dmapOut.println(getAesCiphertext("show " + id));
            response = getAesPlaintext(dmapIn.readLine());
            if (response.startsWith("error")) {
                this.shell.out().println(response);
                return;
            }
            String[] commands = response.split("\n");
            if (commands.length != 6) {
                logger.severe("error expected different message format");
                return;
            }
            msg.setFrom(new Email(commands[0].split("\\s+")[1]));

            String[] addresses = commands[1].split("\\s+")[1].split(",");
            for (String email : addresses) {
                msg.addTo(new Email(email));
            }

            String[] split = commands[2].split("\\s+", 2);
            msg.setSubject(split.length == 2 ? split[1] : "");

            split = commands[3].split("\\s+", 2);
            msg.setData(split.length == 2 ? split[1] : "");

            split = commands[4].split("\\s+", 2);
            msg.setHash(split.length == 2 ? split[1] : "");

            if (!commands[5].startsWith("ok")) {
                logger.severe("error expected 'ok', got " + commands[5] + " instead");
            }

            String calculatedHash = msg.calculateHash();
            logger.info("Calculated hash: " + calculatedHash);
            if (msg.getHash().equals(calculatedHash)) {
                this.shell.out().println("ok");
            } else {
                this.shell.out().println("error");
            }
        } catch (IllegalBlockSizeException | BadPaddingException | NoSuchAlgorithmException | InvalidKeyException e) {
            logger.severe("Error during encryption/decryption");
            e.printStackTrace();
            shutdown();
        } catch (IOException e) {
            logger.severe("Cannot communicate with server");
            e.printStackTrace();
            shutdown();
        } catch (MalformedInputException e) {
            logger.severe("Error during reading of message");
            e.printStackTrace();
        }
    }

    @Command
    @Override
    public void msg(String to, String subject, String data) {
        logger.info("Received 'msg' command:\tto: " + to + "\tsubject: " + subject + "\tdata: " + 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();
        }
    }

    @Command
    @Override
    public void shutdown() {
        logger.info("Received 'shutdown' command");
        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();
        throw new StopShellException();
    }

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