package org.xbib.net.mail.util;

import java.io.IOException;
import java.io.OutputStream;
import java.net.Socket;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * An OutputStream that wraps the Socket's OutputStream and uses
 * the ScheduledExecutorService to schedule a task to close the
 * socket (aborting the write) if the timeout expires.
 */
public class TimeoutOutputStream extends OutputStream {
    private static final String WRITE_TIMEOUT_MESSAGE = "Write timed out";
    private static final String CANNOT_GET_TIMEOUT_TASK_RESULT_MESSAGE = "Couldn't get result of timeout task";

    private final OutputStream os;
    private final ScheduledExecutorService ses;
    private final Callable<String> timeoutTask;
    private final int timeout;
    private final Socket socket;
    private byte[] b1;
    // Implement timeout with a scheduled task
    private ScheduledFuture<String> sf = null;

    public TimeoutOutputStream(Socket socket, ScheduledExecutorService ses, int timeout) throws IOException {
        this.os = socket.getOutputStream();
        this.ses = ses;
        this.timeout = timeout;
        this.socket = socket;
        timeoutTask = new Callable<String>() {
            @Override
            public String call() throws Exception {
                try {
                    os.close();    // close the stream to abort the write
                } catch (Throwable t) {
                    return t.toString();
                }
                return WRITE_TIMEOUT_MESSAGE;
            }
        };
    }

    @Override
    public synchronized void write(int b) throws IOException {
        if (b1 == null)
            b1 = new byte[1];
        b1[0] = (byte) b;
        this.write(b1);
    }

    @Override
    public synchronized void write(byte[] bs, int off, int len)
            throws IOException {
        if ((off < 0) || (off > bs.length) || (len < 0) ||
                ((off + len) > bs.length) || ((off + len) < 0)) {
            throw new IndexOutOfBoundsException();
        } else if (len == 0) {
            return;
        }

        try {
            try {
                if (timeout > 0)
                    sf = ses.schedule(timeoutTask,
                            timeout, TimeUnit.MILLISECONDS);
            } catch (RejectedExecutionException ex) {
                if (!socket.isClosed()) {
                    throw new IOException("Write aborted due to timeout not enforced", ex);
                }
            }

            try {
                os.write(bs, off, len);
            } catch (IOException e) {
                if (sf != null && !sf.cancel(true)) {
                    throw new IOException(handleTimeoutTaskResult(sf), e);
                }
                throw e;
            }
        } finally {
            if (sf != null)
                sf.cancel(true);
        }
    }

    @Override
    public void close() throws IOException {
        os.close();
        if (sf != null) {
            sf.cancel(true);
        }
    }

    private String handleTimeoutTaskResult(ScheduledFuture<String> sf) {
        boolean wasInterrupted = Thread.interrupted();
        String exceptionMessage = null;
        try {
            return sf.get(timeout, TimeUnit.MILLISECONDS);
        } catch (TimeoutException e) {
            exceptionMessage = String.format("%s %s", e, ses.toString());
        } catch (InterruptedException e) {
            wasInterrupted = true;
            exceptionMessage = e.toString();
        } catch (ExecutionException e) {
            exceptionMessage = e.getCause() == null ? e.toString() : e.getCause().toString();
        } catch (Exception e) {
            exceptionMessage = e.toString();
        } finally {
            if (wasInterrupted) {
                Thread.currentThread().interrupt();
            }
        }

        return String.format("%s. %s", CANNOT_GET_TIMEOUT_TASK_RESULT_MESSAGE, exceptionMessage);
    }
}
