/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.sshd.common.io.nio2;

import java.io.IOException;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.ClosedChannelException;
import java.util.HashMap;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.LinkedTransferQueue;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;

import org.apache.sshd.common.SshException;
import org.apache.sshd.common.future.CloseFuture;
import org.apache.sshd.common.future.DefaultSshFuture;
import org.apache.sshd.common.future.SshFuture;
import org.apache.sshd.common.io.IoHandler;
import org.apache.sshd.common.io.IoService;
import org.apache.sshd.common.io.IoSession;
import org.apache.sshd.common.io.IoWriteFuture;
import org.apache.sshd.common.util.Buffer;
import org.apache.sshd.common.util.CloseableUtils;
import org.apache.sshd.common.util.Readable;

/**
 */
public class Nio2Session extends CloseableUtils.AbstractCloseable implements IoSession {

    private static final AtomicLong sessionIdGenerator = new AtomicLong(100);

    private final long id = sessionIdGenerator.incrementAndGet();
    private final Nio2Service service;
    private final IoHandler handler;
    private final AsynchronousSocketChannel socket;
    private final Map<Object, Object> attributes = new HashMap<Object, Object>();
    private final SocketAddress localAddress;
    private final SocketAddress remoteAddress;

    private final Queue<DefaultIoWriteFuture> writes = new LinkedTransferQueue<DefaultIoWriteFuture>();
    private final AtomicReference<DefaultIoWriteFuture> currentWrite = new AtomicReference<DefaultIoWriteFuture>();

    public Nio2Session(Nio2Service service, IoHandler handler, AsynchronousSocketChannel socket) throws IOException {
        this.service = service;
        this.handler = handler;
        this.socket = socket;
        this.localAddress = socket.getLocalAddress();
        this.remoteAddress = socket.getRemoteAddress();
        log.debug("Creating IoSession on {} from {}", localAddress, remoteAddress);
    }

    public long getId() {
        return id;
    }

    public Object getAttribute(Object key) {
        return attributes.get(key);
    }

    public Object setAttribute(Object key, Object value) {
        return attributes.put(key, value);
    }

    public SocketAddress getRemoteAddress() {
        return remoteAddress;
    }

    public SocketAddress getLocalAddress() {
        return localAddress;
    }

    public void suspend() {
        try {
            this.socket.shutdownInput();
        } catch (IOException e) {
            // Ignore
        }
        try {
            this.socket.shutdownOutput();
        } catch (IOException e) {
            // Ignore
        }
    }

    public IoWriteFuture write(Buffer buffer) {
        log.debug("Writing {} bytes", buffer.available());
        ByteBuffer buf = ByteBuffer.wrap(buffer.array(), buffer.rpos(), buffer.available());
        final DefaultIoWriteFuture future = new DefaultIoWriteFuture(null, buf);
        if (state.get() != OPENED) {
            Throwable exc = new ClosedChannelException();
            future.setException(exc);
            exceptionCaught(exc);
            return future;
        }
        writes.add(future);
        startWriting();
        return future;
    }

    private void exceptionCaught(Throwable exc) {
        if (!closeFuture.isClosed()) {
            if (isClosing() || !socket.isOpen()) {
                close(true);
            } else {
                try {
                    log.debug("Caught exception, now calling handler");
                    handler.exceptionCaught(this, exc);
                } catch (Throwable t) {
                    log.info("Exception handler threw exception, closing the session", t);
                    close(true);
                }
            }
        }
    }

    @Override
    protected CloseFuture doCloseGracefully() {
        synchronized (writes) {
            return CloseableUtils.parallel(writes.toArray(new SshFuture[writes.size()]));
        }
    }

    @Override
    protected void doCloseImmediately() {
        for (;;) {
            DefaultIoWriteFuture future = writes.poll();
            if (future != null) {
                future.setException(new ClosedChannelException());
            } else {
                break;
            }
        }
        try {
            socket.close();
        } catch (IOException e) {
            log.info("Exception caught while closing socket", e);
        }
        service.sessionClosed(this);
        super.doCloseImmediately();
        try {
            handler.sessionClosed(this);
        } catch (Exception e) {
            // Ignore
            log.debug("Exception caught while calling IoHandler#sessionClosed", e);
        }
    }

    public IoService getService() {
        return service;
    }

    public void startReading() {
        final ByteBuffer buffer = ByteBuffer.allocate(32 * 1024);
        socket.read(buffer, null, new Nio2CompletionHandler<Integer, Object>() {
            protected void onCompleted(Integer result, Object attachment) {
                try {
                    if (result >= 0) {
                        log.debug("Read {} bytes", result);
                        buffer.flip();
                        Readable buf = new Readable() {
                            public int available() {
                                return buffer.remaining();
                            }
                            public void getRawBytes(byte[] data, int offset, int len) {
                                buffer.get(data, offset, len);
                            }
                        };
                        handler.messageReceived(Nio2Session.this, buf);
                        if (!closeFuture.isClosed()) {
                            startReading();
                        } else {
                            log.debug("IoSession has been closed, stop reading");
                        }
                    } else {
                        log.debug("Socket has been disconnected, closing IoSession now");
                        Nio2Session.this.close(true);
                    }
                } catch (Throwable exc) {
                    failed(exc, attachment);
                }
            }
            protected void onFailed(Throwable exc, Object attachment) {
                exceptionCaught(exc);
            }
        });
    }

    private void startWriting() {
        final DefaultIoWriteFuture future = writes.peek();
        if (future != null) {
            if (currentWrite.compareAndSet(null, future)) {
                try {
                    socket.write(future.buffer, null, new Nio2CompletionHandler<Integer, Object>() {
                        protected void onCompleted(Integer result, Object attachment) {
                            if (future.buffer.hasRemaining()) {
                                try {
                                    socket.write(future.buffer, null, this);
                                } catch (Throwable t) {
                                    log.debug("Exception caught while writing", t);
                                    future.setWritten();
                                    finishWrite();
                                }
                            } else {
                                log.debug("Finished writing");
                                future.setWritten();
                                finishWrite();
                            }
                        }
                        protected void onFailed(Throwable exc, Object attachment) {
                            future.setException(exc);
                            exceptionCaught(exc);
                            finishWrite();
                        }
                        private void finishWrite() {
                            synchronized (writes) {
                                writes.remove(future);
                                writes.notifyAll();
                            }
                            currentWrite.compareAndSet(future, null);
                            startWriting();
                        }
                    });
                } catch (RuntimeException e) {
                    future.setWritten();
                    throw e;
                }
            }
        }
    }

    static class DefaultIoWriteFuture extends DefaultSshFuture<IoWriteFuture> implements IoWriteFuture {
        private final ByteBuffer buffer;
        DefaultIoWriteFuture(Object lock, ByteBuffer buffer) {
            super(lock);
            this.buffer = buffer;
        }
        public void verify() throws SshException {
            try {
                await();
            }
            catch (InterruptedException e) {
                throw new SshException("Interrupted", e);
            }
            if (!isWritten()) {
                throw new SshException("Write failed", getException());
            }
        }

        public boolean isWritten() {
            return getValue() instanceof Boolean;
        }
        public void setWritten() {
            setValue(Boolean.TRUE);
        }
        public Throwable getException() {
            Object v = getValue();
            return v instanceof Throwable ? (Throwable) v : null;
        }
        public void setException(Throwable exception) {
            if (exception == null) {
                throw new IllegalArgumentException("exception");
            }
            setValue(exception);
        }
    }

    public String toString() {
        return getClass().getSimpleName() + "[local=" + localAddress + ", remote=" + remoteAddress + "]";
    }
}
