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

import com.jscape.inet.ssh.protocol.messages.IdentificationString;
import com.jscape.inet.ssh.protocol.messages.Message;
import com.jscape.inet.ssh.protocol.v2.authentication.ClientAuthentication;
import com.jscape.inet.ssh.protocol.v2.marshaling.ClientSession;
import com.jscape.inet.ssh.protocol.v2.marshaling.IdentificationStringCodec;
import com.jscape.inet.ssh.protocol.v2.marshaling.MessageCodec;
import com.jscape.inet.ssh.protocol.v2.marshaling.algorithms.CompressionFactory;
import com.jscape.inet.ssh.protocol.v2.marshaling.algorithms.EncryptionFactory;
import com.jscape.inet.ssh.protocol.v2.marshaling.algorithms.MacFactory;
import com.jscape.inet.ssh.protocol.v2.messages.*;
import com.jscape.inet.ssh.protocol.v2.transport.TransportConnection;
import com.jscape.inet.ssh.protocol.v2.transport.TransportMessages;
import com.jscape.inet.ssh.protocol.v2.transport.keyexchange.ClientKeyExchangeFactory;
import com.jscape.inet.ssh.protocol.v2.transport.keyexchange.HostKeyVerificationService;
import com.jscape.util.SuppressObfuscation;
import com.jscape.util.io.DataCodec;
import com.jscape.util.net.connection.*;

import java.util.logging.Level;
import java.util.logging.Logger;

import static com.jscape.util.Parameters.assertGreater;
import static com.jscape.util.Parameters.assertNotNull;

/**
 * @author Alex Usun
 */
public class SessionConnector
        implements Connector<SessionConnection> {

    private static final DataCodec<Message> INITIAL_CODEC = new IdentificationStringCodec();
    private static final String LANGUAGE_TAG = "en";

    private final Connector<SyncRawConnection> base;
    private final IdentificationString identificationString;
    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 ClientAuthentication authentication;
    private final int initialWindowSize;
    private final int maxPacketSize;
    private final Logger logger;

    public SessionConnector(Connector<SyncRawConnection> base,
                            IdentificationString identificationString,
                            PublicKeyAlgorithm[] publicKeyAlgorithms,
                            ClientKeyExchangeFactory keyExchangeFactory,
                            EncryptionFactory encryptionFactory,
                            MacFactory macFactory,
                            CompressionFactory compressionFactory,
                            HostKeyVerificationService hostKeyVerificationService,
                            ClientAuthentication authentication,
                            int initialWindowSize,
                            int maxPacketSize,
                            Logger logger) {
        assertNotNull(base);
        this.base = base;

        assertNotNull(identificationString);
        this.identificationString = identificationString;

        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;

        assertNotNull(authentication);
        this.authentication = authentication;

        assertGreater(initialWindowSize, 0, "Bad initial window size value.");
        this.initialWindowSize = initialWindowSize;

        assertGreater(maxPacketSize, 0, "Bad max packet size value.");
        this.maxPacketSize = maxPacketSize;

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

    @Override
    public SessionConnection connect()
            throws OperationException {
        SyncRawConnection baseConnection = this.base.connect();
        SyncMessageConnectionSyncRawBase<Message> messageConnection = new SyncMessageConnectionSyncRawBase<>(baseConnection, INITIAL_CODEC);
        logConnectionCreated(messageConnection);

        try {
            ClientSession session = exchangeProtocolVersion(messageConnection);
            TransportConnection transportConnection = new TransportConnection(messageConnection,
                                                                              session,
                                                                              this.publicKeyAlgorithms,
                                                                              this.keyExchangeFactory,
                                                                              this.encryptionFactory,
                                                                              this.macFactory,
                                                                              this.compressionFactory,
                                                                              this.hostKeyVerificationService,
                                                                              this.logger);
            transportConnection.exchangeKeys();
            requestAuthenticationService(transportConnection);
            this.authentication.applyTo(transportConnection, ServiceType.CONNECTION.name);

            SessionConnectionMessages.init(session.messageCodec);

            return new SessionConnection(transportConnection, this.initialWindowSize, this.maxPacketSize);
        } catch (UnsupportedServerException e) {
            disconnect(DisconnectReason.PROTOCOL_VERSION_NOT_SUPPORTED, String.format("Unsupported protocol version: %s.", e.getMessage()), messageConnection);
            throw new OperationException(String.format("Unsupported server protocol version: %s.", e.getMessage()));
        } catch (Exception e) {
            baseConnection.close();
            throw new OperationException(e);
        }
    }

    private ClientSession exchangeProtocolVersion(SyncMessageConnectionSyncRawBase<Message> connection)
            throws Exception {
        send(this.identificationString, connection);
        IdentificationString serverIdentificationString = receive(connection);

        ClientSession session = new ClientSession(TransportMessages.init(new MessageCodec()), this.identificationString, serverIdentificationString);
        connection.messageCodec(session.packetCodec);

        assertServerSupported(serverIdentificationString);

        return session;
    }

    private void assertServerSupported(IdentificationString serverIdentificationString)
            throws UnsupportedServerException {
        String serverProtocolVersion = "1.99".equals(serverIdentificationString.protocolVersion) ? "2.0" : serverIdentificationString.protocolVersion;

        if (!this.identificationString.protocolVersion.equals(serverProtocolVersion)) {
            throw new UnsupportedServerException(serverIdentificationString.protocolVersion);
        }
    }

    private void requestAuthenticationService(TransportConnection transportConnection)
            throws Exception {
        transportConnection.write(new SshMsgServiceRequest(ServiceType.USERAUTH));
        transportConnection.<SshMsgServiceAccept>read();
    }

    private void disconnect(DisconnectReason reason, String message, MessageConnection<Message> connection) {
        try {
            send(new SshMsgDisconnect(reason, message, LANGUAGE_TAG), connection);
        } catch (Exception e) {
            logError(e);
        } finally {
            close(connection);
        }
    }

    private void send(Message message, MessageConnection<Message> connection)
            throws Exception {
        logOutgoingMessage(connection, message);
        connection.write(message);
    }

    @SuppressWarnings("unchecked")
    private <M extends Message> M receive(SyncMessageConnectionSyncRawBase<Message> connection)
            throws Exception {
        Message message = connection.read();
        logIncomingMessage(connection, message);

        return (M) message;
    }

    private void close(MessageConnection<Message> connection) {
        if (!connection.closed()) {
            connection.close();
        }
    }

    private void logConnectionCreated(Connection connection) {
        if (this.logger.isLoggable(Level.FINE)) {
            this.logger.log(Level.FINE, String.format("SSH connection created: [%s <-> %s].",
                                                      connection.localAddress(),
                                                      connection.remoteAddress()));
        }
    }

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

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

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

    @SuppressObfuscation
    private static class UnsupportedServerException
            extends Exception {

        private UnsupportedServerException(String serverVersion) {
            super(serverVersion);
        }
    }
}
