package com.jscape.inet.ssh.protocol.v2.transport;

import com.jscape.inet.ssh.protocol.messages.Message;
import com.jscape.inet.ssh.protocol.v2.marshaling.ClientSession;
import com.jscape.inet.ssh.protocol.v2.marshaling.algorithms.*;
import com.jscape.inet.ssh.protocol.v2.messages.*;
import com.jscape.inet.ssh.protocol.v2.transport.keyexchange.ClientKeyExchange;
import com.jscape.inet.ssh.protocol.v2.transport.keyexchange.ClientKeyExchangeFactory;
import com.jscape.inet.ssh.protocol.v2.transport.keyexchange.HostKeyVerificationService;
import com.jscape.inet.ssh.protocol.v2.transport.keyexchange.KeyExchangeResult;
import com.jscape.util.Clock;
import com.jscape.util.SuppressObfuscation;
import com.jscape.util.net.TransportAddress;
import com.jscape.util.net.connection.MessageConnection;
import com.jscape.util.net.connection.SyncMessageConnection;

import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Collectors;

import static com.jscape.util.Parameters.assertNotNull;
import static com.jscape.util.net.connection.Connection.ConnectionException.wrap;

/**
 * @author Alex Usun
 */
public class TransportConnection
        implements SyncMessageConnection<Message>,
                   SshMsgDisconnect.Handler,
                   SshMsgIgnore.Handler,
                   SshMsgDebug.Handler,
                   SshMsgKexInit.Handler {

    private static final String LANGUAGE_TAG = "en";

    private static final List<String> LANGUAGES = Collections.emptyList();

    private final SyncMessageConnection<Message> base;
    private final ClientSession session;
    private final PublicKeyAlgorithm[] publicKeyAlgorithms;
    private final ClientKeyExchangeFactory keyExchangeFactory;
    private final EncryptionFactory encryptionFactory;
    private final MacFactory macFactory;
    private final CompressionFactory compressionFactory;
    private final HostKeyVerificationService hostKeyVerificationService;
    private final Lock sendLock;
    private final Lock keyExchangeLock;
    private final AtomicBoolean closed;
    private final Logger logger;
    private boolean keyExchangeStarted;
    private boolean keyExchangeInProgress;
    private Long lastKeyExchangeDate;
    private Exception error;

    public TransportConnection(SyncMessageConnection<Message> base,
                               ClientSession session,
                               PublicKeyAlgorithm[] publicKeyAlgorithms,
                               ClientKeyExchangeFactory keyExchangeFactory,
                               EncryptionFactory encryptionFactory,
                               MacFactory macFactory,
                               CompressionFactory compressionFactory,
                               HostKeyVerificationService hostKeyVerificationService,
                               Logger logger) {
        assertNotNull(base);
        this.base = base;

        assertNotNull(session);
        this.session = session;

        assertNotNull(publicKeyAlgorithms);
        this.publicKeyAlgorithms = publicKeyAlgorithms;

        assertNotNull(keyExchangeFactory);
        this.keyExchangeFactory = keyExchangeFactory;

        assertNotNull(encryptionFactory);
        this.encryptionFactory = encryptionFactory;

        assertNotNull(macFactory);
        this.macFactory = macFactory;

        assertNotNull(compressionFactory);
        this.compressionFactory = compressionFactory;

        assertNotNull(hostKeyVerificationService);
        this.hostKeyVerificationService = hostKeyVerificationService;

        this.sendLock = new ReentrantLock();
        this.keyExchangeLock = new ReentrantLock();
        this.closed = new AtomicBoolean(false);

        assertNotNull(logger);
        this.logger = logger;
    }

    @Override
    public long creationTime() {
        return this.base.creationTime();
    }

    @Override
    public TransportAddress localAddress() {
        return this.base.localAddress();
    }

    @Override
    public TransportAddress remoteAddress() {
        return this.base.remoteAddress();
    }

    @Override
    public Map<Object, Object> attributes() {
        return this.base.attributes();
    }

    @Override
    public boolean closed() {
        return this.closed.get();
    }

    @Override
    public void close() {
        if (this.closed.getAndSet(true)) {
            return;
        }

        this.base.close();
    }

    @Override
    public Message read()
            throws ConnectionException {
        Message message;

        do {
            message = receive();
        } while (message == null);

        return message;
    }

    @Override
    public void write(Message message)
            throws ConnectionException {
        this.keyExchangeLock.lock();

        try {
            send(message);
        } finally {
            this.keyExchangeLock.unlock();
        }
    }

    @Override
    public void handle(SshMsgDisconnect message, MessageConnection<Message> connection) {
        close();
    }

    @Override
    public void handle(SshMsgIgnore message, MessageConnection<Message> connection) {
        // empty
    }

    @Override
    public void handle(SshMsgDebug message, MessageConnection<Message> connection) {
        // empty
    }

    @Override
    public void handle(SshMsgKexInit message, MessageConnection<Message> connection) {
        this.keyExchangeLock.lock();

        try {
            assertNoKeyExchange();

            this.session.serverKexInitMessage = message;

            if (!this.keyExchangeStarted) {
                sendKexInitMessage();
            }

            this.keyExchangeStarted = false;
            this.keyExchangeInProgress = true;

            this.session.initAlgorithms();
            logSessionAlgorithms();

            ClientKeyExchange keyExchange = this.keyExchangeFactory.keyExchangeFor(this.session.algorithms.keyExchange);
            KeyExchangeResult keyExchangeResult = keyExchange.exchangeKeys(this, this.session, this.hostKeyVerificationService);

            send(new SshMsgNewKeys());

            this.session.hostKey = keyExchangeResult.hostKey;
            this.session.initSessionId(keyExchangeResult.exchangeHash);
            initWritingAlgorithms(keyExchangeResult);

            SshMsgNewKeys newKeysMessage = (SshMsgNewKeys) read();
            initReadingAlgorithms(keyExchangeResult);

            this.keyExchangeInProgress = false;
            this.lastKeyExchangeDate = Clock.now();
        } catch (ProtocolException e) {
            this.error = e;
            disconnect(DisconnectReason.PROTOCOL_ERROR, "Unexpected SSH_MSG_KEX_INIT message.");
        } catch (Algorithms.CommonAlgorithmsNotFoundException e) {
            logError(e);
            this.error = e;
            disconnect(DisconnectReason.KEY_EXCHANGE_FAILED, "No common algorithms found.");
        } catch (Exception e) {
            logError(e);
            this.error = e;
            disconnect(DisconnectReason.BY_APPLICATION, "Internal server error.");
        } finally {
            this.keyExchangeLock.unlock();
        }
    }

    public SyncMessageConnection<Message> base() {
        return this.base;
    }

    public void exchangeKeys()
            throws ConnectionException {
        this.keyExchangeLock.lock();

        try {
            sendKexInitMessage();
            this.keyExchangeStarted = true;

            while (this.keyExchangeStarted || this.keyExchangeInProgress) {
                receive();
            }
        } finally {
            this.keyExchangeLock.unlock();
        }
    }

    public void debug(boolean alwaysDisplay, String message, String languageTag)
            throws ConnectionException {
        send(new SshMsgDebug(alwaysDisplay, message, languageTag));
    }

    public void ignore(byte[] data)
            throws ConnectionException {
        send(new SshMsgIgnore(data));
    }

    public void disconnect(DisconnectReason reason, String message) {
        try {
            send(new SshMsgDisconnect(reason, message, LANGUAGE_TAG));
        } catch (ConnectionException e) {
            logErrorSendMessage(e);
        } finally {
            close();
        }
    }

    public ClientSession session() {
        return this.session;
    }

    public long readBytes() {
        return this.session.packetCodec.readBytes();
    }

    public long writtenBytes() {
        return this.session.packetCodec.writtenBytes();
    }

    public long readMessages() {
        return this.session.packetCodec.readMessages();
    }

    public long writtenMessages() {
        return this.session.packetCodec.writtenMessages();
    }

    public Long lastKeyExchangeDate() {
        return this.lastKeyExchangeDate;
    }

    @Override
    public String toString() {
        return this.base.toString();
    }

    private void send(Message message)
            throws ConnectionException {
        this.sendLock.lock();

        try {
            logOutgoingMessage(message);
            this.base.write(message);
        } finally {
            this.sendLock.unlock();
        }
    }

    private Message receive()
            throws ConnectionException {
        Message message = this.base.read();
        logIncomingMessage(message);

        if (isTransportMessage(message)) {
            dispatch(message);

            if (closed() && this.error != null) {
                throw wrap(this.error);
            }

            return null;
        } else {
            return message;
        }
    }

    private void sendKexInitMessage()
            throws ConnectionException {
        List<String> keyExchanges = this.keyExchangeFactory.keyExchanges();
        List<String> encryptions = this.encryptionFactory.encryptions();
        List<String> macs = this.macFactory.macs();
        List<String> compressions = this.compressionFactory.compressions();

        this.session.clientKexInitMessage = SshMsgKexInit.messageFor(keyExchanges,
                                                                     Arrays.stream(this.publicKeyAlgorithms).map(a -> a.identifier).collect(Collectors.toList()),
                                                                     encryptions,
                                                                     encryptions,
                                                                     macs,
                                                                     macs,
                                                                     compressions,
                                                                     compressions,
                                                                     LANGUAGES,
                                                                     LANGUAGES);


        send(this.session.clientKexInitMessage);
    }

    @SuppressWarnings("unchecked")
    private boolean isTransportMessage(Message message) {
        return message.handlerClass().isAssignableFrom(getClass());
    }

    @SuppressWarnings("unchecked")
    private void dispatch(Message message)
            throws ConnectionException {
        try {
            message.accept(this, this.base);
        } catch (Exception e) {
            logUnsupportedMessage(message);
            throw new ConnectionException("Unsupported incoming message.", e);
        }
    }

    private void initReadingAlgorithms(KeyExchangeResult keyExchangeResult)
            throws Exception {
        KeyFactory keyFactory = keyExchangeResult.asKeyFactory(this.session.sessionId);

        Encryption encryption = this.encryptionFactory.serverToClientEncryptionFor(this.session.algorithms.encryptionServerToClient, keyFactory, EncryptionFactory.Mode.DECRYPTION);
        Mac mac = this.macFactory.serverToClientMacFor(this.session.algorithms.macServerToClient, keyFactory);
        Compression compression = this.compressionFactory.serverToClientCompressionFor(this.session.algorithms.compressionServerToClient, CompressionFactory.Mode.DECOMPRESSION);

        this.session.readingAlgorithms(encryption, mac, compression);
    }

    private void initWritingAlgorithms(KeyExchangeResult keyExchangeResult)
            throws Exception {
        KeyFactory keyFactory = keyExchangeResult.asKeyFactory(this.session.sessionId);

        Encryption encryption = this.encryptionFactory.clientToServerEncryptionFor(this.session.algorithms.encryptionClientToServer, keyFactory, EncryptionFactory.Mode.ENCRYPTION);
        Mac mac = this.macFactory.clientToServerMacFor(this.session.algorithms.macClientToServer, keyFactory);
        Compression compression = this.compressionFactory.clientToServerCompressionFor(this.session.algorithms.compressionClientToServer, CompressionFactory.Mode.COMPRESSION);

        this.session.writingAlgorithms(encryption, mac, compression);
    }

    private void assertNoKeyExchange()
            throws ProtocolException {
        if (this.keyExchangeInProgress) {
            throw new ProtocolException();
        }
    }

    private void logIncomingMessage(Message message) {
        if (this.logger.isLoggable(Level.FINER)) {
            this.logger.log(Level.FINER,
                            String.format("Incoming SSH message at [%s <-> %s]: %s",
                                          localAddress(),
                                          remoteAddress(),
                                          message));
        }
    }

    private void logUnsupportedMessage(Message message) {
        if (this.logger.isLoggable(Level.WARNING)) {
            this.logger.log(Level.WARNING,
                            String.format("Unsupported SSH message at [%s <-> %s]: %s",
                                          localAddress(),
                                          remoteAddress(),
                                          message));
        }
    }

    private void logOutgoingMessage(Message message) {
        if (this.logger.isLoggable(Level.FINER)) {
            this.logger.log(Level.FINER,
                            String.format("Outgoing SSH message at [%s <-> %s]: %s",
                                          localAddress(),
                                          remoteAddress(),
                                          message));
        }
    }

    private void logErrorSendMessage(Throwable error) {
        if (this.logger.isLoggable(Level.WARNING) && !this.base.closed()) {
            this.logger.log(Level.WARNING, "Error sending message.", error);
        }
    }

    private void logSessionAlgorithms() {
        if (this.logger.isLoggable(Level.FINER)) {
            this.logger.log(Level.FINER,
                            String.format("SSH session algorithms at [%s <-> %s]: %s",
                                          this.base.localAddress(),
                                          this.base.remoteAddress(),
                                          this.session.algorithms));
        }
    }

    private void logError(Throwable error) {
        if (this.logger.isLoggable(Level.WARNING)) {
            this.logger.log(Level.WARNING, "Operation error.", error);
        }
    }

    @SuppressObfuscation
    private static class ProtocolException
            extends Exception {
    }

}
