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

import com.jscape.inet.ssh.protocol.messages.Message;
import com.jscape.inet.ssh.protocol.v2.messages.*;
import com.jscape.inet.ssh.protocol.v2.transport.TransportConnection;
import com.jscape.util.net.TransportAddress;
import com.jscape.util.net.connection.Connection;
import com.jscape.util.net.connection.MessageConnection;

import java.security.PublicKey;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
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 SessionConnection
        implements Connection,
                   SshMsgGlobalRequestUnknown.Handler {

    private final TransportConnection base;
    private final int initialWindowSize;
    private final int maxPacketSize;
    private final AtomicBoolean closed;
    private final Logger logger;
    private int channelIds;
    private ChannelConnection currentChannel;

    public SessionConnection(TransportConnection base, int initialWindowSize, int maxPacketSize) {
        assertNotNull(base);
        this.base = base;

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

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

        this.closed = new AtomicBoolean(false);

        this.logger = Logger.getLogger("com.jscape.inet.ssh");
    }

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

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

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

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

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

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

        this.base.close();
    }

    public TransportConnection base() {
        return this.base;
    }

    public PublicKey hostKey() {
        return this.base.session().hostKey;
    }

    @Override
    public void handle(SshMsgGlobalRequestUnknown message, MessageConnection<Message> connection) {
        if (message.wantReply) {
            sendSafe(new SshMsgRequestFailure());
        }
    }

    public SessionChannelConnection openSessionChannel(boolean closeChannelConfirmationRequired)
            throws ConnectionException {
        assertNotBusy();

        int localId = nextChannelId();
        this.base.write(new SshMsgChannelOpenSession(localId, this.initialWindowSize, this.maxPacketSize));
        Message response = read();
        assertChannelOpenSuccess(response);

        SshMsgChannelOpenConfirmation confirmation = (SshMsgChannelOpenConfirmation) response;

        this.currentChannel = new SessionChannelConnection(this.base,
                                                           localId,
                                                           confirmation.senderChannel,
                                                           this.initialWindowSize,
                                                           this.maxPacketSize,
                                                           confirmation.initialWindowSize,
                                                           confirmation.maxPacketSize,
                                                           closeChannelConfirmationRequired);

        return (SessionChannelConnection) this.currentChannel;
    }

    public ChannelConnection openDirectTcpIpChannel(String host, int port, String originatorAddress, int originatorPort, boolean closeChannelConfirmationRequired)
            throws ConnectionException {
        assertNotBusy();

        int localId = nextChannelId();
        this.base.write(new SshMsgChannelOpenDirectTcpIp(localId, this.initialWindowSize, this.maxPacketSize, host, port, originatorAddress, originatorPort));
        Message response = read();
        assertChannelOpenSuccess(response);

        SshMsgChannelOpenConfirmation confirmation = (SshMsgChannelOpenConfirmation) response;

        this.currentChannel = new ChannelConnection(this.base,
                                                    localId,
                                                    confirmation.senderChannel,
                                                    this.initialWindowSize,
                                                    this.maxPacketSize,
                                                    confirmation.initialWindowSize,
                                                    confirmation.maxPacketSize,
                                                    closeChannelConfirmationRequired);

        return this.currentChannel;
    }

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

    private void sendSafe(Message message) {
        try {
            this.base.write(message);
        } catch (Exception e) {
            logErrorSendMessage(e);
        }
    }

    private Message read()
            throws ConnectionException {
        Message message;

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

        return message;
    }

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

        if (dispatchable(message)) {
            dispatch(message);
            return null;
        } else {
            return message;
        }
    }

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

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

    private void assertNotBusy()
            throws ConnectionException {
        if (this.currentChannel != null && !this.currentChannel.closed()) {
            throw new ConnectionException("Multiple channels are not supported.");
        }
    }

    private int nextChannelId() {
        if (this.channelIds == Integer.MAX_VALUE) {
            this.channelIds = 0;
        }

        return ++this.channelIds;
    }

    private void assertChannelOpenSuccess(Message response)
            throws ConnectionException {
        if (!(response instanceof SshMsgChannelOpenConfirmation)) {
            throw new ConnectionException("Error opening session channel.");
        }
    }

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

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