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.util.Time;
import com.jscape.util.io.ByteArrayOutputStream;
import com.jscape.util.net.TransportAddress;
import com.jscape.util.net.connection.Connection.ConnectionClosedException;
import com.jscape.util.net.connection.Connection.ConnectionException;
import com.jscape.util.net.connection.MessageConnection;

import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
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 Channel
        implements SshMsgChannelWindowAdjust.Handler,
                   SshMsgChannelData.Handler,
                   SshMsgChannelExtendedData.Handler,
                   SshMsgChannelEof.Handler,
                   SshMsgChannelClose.Handler {

    protected static final Time WRITE_ALLOWED_TIMEOUT = Time.millis(1);

    protected final MessageConnection<Message> connection;
    protected final int localId;
    protected final int remoteId;
    protected final Window localWindow;
    protected final long localWindowMinSize;
    protected final int localMaxPacketSize;
    protected final Window remoteWindow;
    protected final int remoteMaxPacketSize;
    protected final Set<Listener> listeners;
    protected final Logger logger;
    protected Writer writer;
    protected volatile boolean closed;

    public Channel(MessageConnection<Message> connection,
                   int localId,
                   int remoteId,
                   long localInitialWindowSize,
                   int localMaxPacketSize,
                   long remoteInitialWindowSize,
                   int remoteMaxPacketSize) {
        assertNotNull(connection);
        this.connection = connection;

        this.localId = localId;
        this.remoteId = remoteId;

        this.localWindow = new Window(localInitialWindowSize);
        this.localWindowMinSize = localInitialWindowSize / 3;

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

        this.remoteWindow = new Window(remoteInitialWindowSize);

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

        this.listeners = new LinkedHashSet<Listener>();

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

        this.writer = new NonBlockingWriter();
    }

    public int localId() {
        return this.localId;
    }

    public int remoteId() {
        return this.remoteId;
    }

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

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

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

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

    public MessageConnection<Message> connection() {
        return this.connection;
    }

    @Override
    public void handle(SshMsgChannelWindowAdjust message, MessageConnection<Message> connection) {
        this.writer.adjustWindow(message.bytesToAdd);
    }

    @Override
    public void handle(SshMsgChannelData message, MessageConnection<Message> connection) {
        raiseDataEvent(message);
        this.localWindow.decrease(message.length);
        adjustLocalWindow();
    }

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

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

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

    public void sendRequestSuccess(boolean wantReply) {
        if (wantReply) {
            send(new SshMsgChannelSuccess(this.remoteId));
        }
    }

    public void sendRequestFailure(boolean wantReply) {
        if (wantReply) {
            send(new SshMsgChannelFailure(this.remoteId));
        }
    }

    public void bind(Listener listener) {
        assertNotNull(listener);
        this.listeners.add(listener);
    }

    public void unbind(Listener listener) {
        this.listeners.remove(listener);
    }

    public void close() {
        if (this.closed) {
            return;
        }

        this.closed = true;
        send(new SshMsgChannelClose(this.remoteId));
        raiseClosedEvent();
        this.writer.close();
    }

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

    public void write(byte[] data, int offset, int length)
            throws ConnectionException {
        write(data, offset, length, false);
    }

    public void writeStderr(byte[] data, int offset, int length)
            throws ConnectionException {
        write(data, offset, length, true);
    }

    public void eof() {
        send(new SshMsgChannelEof(this.remoteId));
    }

    public void blockingWrite() {
        if (this.writer instanceof NonBlockingWriter) {
            this.writer = new BlockingWriter();
        }
    }

    @Override
    public String toString() {
        return "Channel {"
                + "localId=" + this.localId
                + ", remoteId=" + this.remoteId
                + ", localWindow=" + this.localWindow
                + ", localWindowMinSize=" + this.localWindowMinSize
                + ", localMaxPacketSize=" + this.localMaxPacketSize
                + ", remoteWindow=" + this.remoteWindow
                + ", remoteMaxPacketSize=" + this.remoteMaxPacketSize
                + ", writer=" + this.writer
                + ", closed=" + this.closed
                + '}';
    }

    protected void write(byte[] data, int offset, int length, boolean stderr)
            throws ConnectionException {
        while (length > 0 && !this.closed) {
            int written = this.writer.write(data, offset, length, stderr);
            offset += written;
            length -= written;
        }

        if (this.closed) {
            throw new ConnectionClosedException();
        }
    }

    protected void adjustLocalWindow() {
        if (this.localWindow.size() <= this.localWindowMinSize) {
            long bytesToAdd = this.localWindow.space();
            send(new SshMsgChannelWindowAdjust(this.remoteId, bytesToAdd));
            this.localWindow.increase(bytesToAdd);
        }
    }

    protected void send(Message message) {
        try {
            this.connection.write(message);
        } catch (Exception e) {
            logErrorSendingMessage(e);
        }
    }

    protected void sendData(byte[] data, int offset, int length, boolean stderr)
            throws ConnectionException {
        if (stderr) {
            this.connection.write(new SshMsgChannelExtendedData(this.remoteId, DataType.SSH_EXTENDED_DATA_STDERR, data, offset, length));
        } else {
            this.connection.write(new SshMsgChannelData(this.remoteId, data, offset, length));
        }

        this.remoteWindow.decrease(length);
    }

    protected int nextDataPacketSize(int length) {
        return Math.min(this.remoteMaxPacketSize, (int) Math.min(this.remoteWindow.size(), length));
    }

    protected void raiseDataEvent(SshMsgChannelData message) {
        for (Listener listener : this.listeners) {
            listener.onData(this, message.data);
        }
    }

    protected void raiseEofEvent() {
        for (Listener listener : this.listeners) {
            listener.onEof(this);
        }
    }

    protected void raiseClosedEvent() {
        for (Listener listener : this.listeners) {
            listener.onClosed(this);
        }
    }

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

    public interface Listener {

        void onData(Channel channel, byte[] data);

        void onEof(Channel channel);

        void onClosed(Channel channel);
    }

    protected interface Writer {

        int write(byte[] data, int offset, int length, boolean stderr)
                throws ConnectionException;

        void adjustWindow(long size);

        void close();
    }

    protected class BlockingWriter
            implements Writer {

        private final Lock writeLock;
        private final Condition writeAllowed;
        private volatile boolean closed;

        public BlockingWriter() {
            this.writeLock = new ReentrantLock();
            this.writeAllowed = this.writeLock.newCondition();
        }

        @Override
        public int write(byte[] data, int offset, int length, boolean stderr)
                throws ConnectionException {
            this.writeLock.lock();

            try {
                awaitRemoteWindow();
                int packetSize = nextDataPacketSize(length);
                sendData(data, offset, packetSize, stderr);

                return packetSize;
            } finally {
                this.writeLock.unlock();
            }
        }

        @Override
        public void adjustWindow(long size) {
            this.writeLock.lock();

            try {
                remoteWindow.increase(size);
                this.writeAllowed.signalAll();
            } finally {
                this.writeLock.unlock();
            }
        }

        @Override
        public void close() {
            this.writeLock.lock();

            try {
                this.closed = true;
                this.writeAllowed.signalAll();
            } finally {
                this.writeLock.unlock();
            }
        }

        private void awaitRemoteWindow() {
            while (remoteWindow.size() == 0 && !this.closed) {
                try {
                    this.writeAllowed.await(WRITE_ALLOWED_TIMEOUT.value, WRITE_ALLOWED_TIMEOUT.unit);
                } catch (InterruptedException e) {
                    // ignore
                }
            }
        }
    }

    protected class NonBlockingWriter
            implements Writer {

        private final ByteArrayOutputStream dataBuffer;
        private final ByteArrayOutputStream stderrBuffer;
        private final Lock writeLock;

        public NonBlockingWriter() {
            this.dataBuffer = new ByteArrayOutputStream(2 * remoteMaxPacketSize);
            this.stderrBuffer = new ByteArrayOutputStream(2 * remoteMaxPacketSize);
            this.writeLock = new ReentrantLock();
        }

        @Override
        public int write(byte[] data, int offset, int length, boolean stderr)
                throws ConnectionException {
            this.writeLock.lock();

            try {
                int packetSize = nextDataPacketSize(length);
                ByteArrayOutputStream buffer = stderr ? this.stderrBuffer : this.dataBuffer;

                if (packetSize > 0 && buffer.isEmpty()) {
                    sendData(data, offset, packetSize, stderr);
                    return packetSize;
                } else {
                    buffer.write(data, offset, length);
                    return length;
                }
            } finally {
                this.writeLock.unlock();
            }
        }

        @Override
        public void adjustWindow(long size) {
            this.writeLock.lock();

            try {
                remoteWindow.increase(size);
                sendBufferedData(this.dataBuffer);
                sendBufferedData(this.stderrBuffer);
            } finally {
                this.writeLock.unlock();
            }
        }

        @Override
        public void close() {
            // empty
        }

        private void sendBufferedData(ByteArrayOutputStream buffer) {
            try {
                int offset = 0;
                int length = buffer.position();
                int packetSize;
                while (length > 0 && (packetSize = nextDataPacketSize(length)) > 0) {
                    sendData(buffer.buffer(), offset, packetSize, buffer == this.stderrBuffer);
                    offset += packetSize;
                    length -= packetSize;
                }

                buffer.relocate(offset);
            } catch (Exception e) {
                logErrorSendingMessage(e);
            }
        }
    }

}
