package org.jboss.resteasy.plugins.server.servlet;

import java.io.IOException;
import java.io.OutputStream;
import java.util.Comparator;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Queue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.atomic.AtomicBoolean;

import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.ws.rs.core.MultivaluedMap;
import jakarta.ws.rs.core.NewCookie;

import org.jboss.resteasy.core.ResteasyContext;
import org.jboss.resteasy.core.ResteasyContext.CloseableContext;
import org.jboss.resteasy.spi.AsyncOutputStream;
import org.jboss.resteasy.spi.HttpRequest;
import org.jboss.resteasy.spi.HttpResponse;
import org.jboss.resteasy.spi.ResteasyProviderFactory;

/**
 * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
 * @version $Revision: 1 $
 */
public class HttpServletResponseWrapper implements HttpResponse {
    public abstract class AsyncOperation {
        final CompletableFuture<Void> future;
        final OutputStream stream;
        private final long id;

        @Deprecated
        public AsyncOperation(final OutputStream stream) {
            this(stream, new CompletableFuture<>(),
                    (stream instanceof DeferredOutputStream ? ((DeferredOutputStream) stream).getId() : -1));
        }

        private AsyncOperation(final DeferredOutputStream stream) {
            this(stream, new CompletableFuture<>(), stream.getId());
        }

        private AsyncOperation(final OutputStream stream, final CompletableFuture<Void> future, final long id) {
            this.stream = stream;
            this.future = future;
            this.id = id;
        }

        public void work(ServletOutputStream sos) {
            try (CloseableContext c = ResteasyContext.addCloseableContextDataLevel(contextDataMap)) {
                doWork(sos);
            }
        }

        protected abstract void doWork(ServletOutputStream sos);

        protected void requeue(final AsyncOperation op) {
            if (op.future.isDone()) {
                return;
            }
            if (stream instanceof DeferredOutputStream) {
                ((DeferredOutputStream) stream).queue(op);
            }
        }

        protected void queueComplete(final AsyncOperation op) {
            if (op.future.isDone()) {
                return;
            }
            if (stream instanceof DeferredOutputStream) {
                final AsyncOperation requeue = (op instanceof CompletionOperation ? op : new CompletionOperation(op));
                ((DeferredOutputStream) stream).queue(requeue);
            }
        }
    }

    public class WriteOperation extends AsyncOperation {

        private final byte[] bytes;
        private final int offset;
        private final int length;

        @Deprecated
        public WriteOperation(final OutputStream stream, final byte[] bytes, final int offset, final int length) {
            super(stream);
            this.bytes = bytes;
            this.offset = offset;
            this.length = length;
        }

        private WriteOperation(final DeferredOutputStream stream, final byte[] bytes, final int offset, final int length) {
            super(stream);
            this.bytes = bytes;
            this.offset = offset;
            this.length = length;
        }

        @Override
        protected void doWork(ServletOutputStream sos) {
            try {
                // we only are complete if isReady says we're good to write, otherwise
                // we will be complete in the next onWritePossible or onError
                if (sos == null) {
                    stream.write(bytes, offset, length);
                    future.complete(null);
                } else {
                    // Check if the stream is ready and if so write the data
                    if (sos.isReady()) {
                        stream.write(bytes, offset, length);
                        // Recheck before we complete the future as the write above may still be in process
                        if (sos.isReady()) {
                            future.complete(null);
                        } else {
                            queueComplete(this);
                        }
                    } else {
                        // The stream is not ready, requeue ourself
                        requeue(this);
                    }
                }
            } catch (IOException e) {
                future.completeExceptionally(e);
            }
        }

        @Override
        public String toString() {
            return "[write: " + new String(bytes) + "]";
        }
    }

    public class FlushOperation extends AsyncOperation {
        @Deprecated
        public FlushOperation(final OutputStream os) {
            super(os);
        }

        public FlushOperation(final DeferredOutputStream os) {
            super(os);
        }

        @Override
        protected void doWork(ServletOutputStream sos) {
            try {
                // we only are complete if isReady says we're good to write, otherwise
                // we will be complete in the next onWritePossible or onError
                if (sos == null) {
                    stream.flush();
                    future.complete(null);
                } else {
                    // The stream is ready, flush the output
                    if (sos.isReady()) {
                        stream.flush();
                        // Recheck before we complete the future as the flush above may still be in process
                        if (sos.isReady()) {
                            future.complete(null);
                        } else {
                            queueComplete(this);
                        }
                    } else {
                        // The stream is not ready, requeue ourself
                        requeue(this);
                    }
                }
            } catch (IOException e) {
                future.completeExceptionally(e);
            }
        }

        @Override
        public String toString() {
            return "[flush]";
        }
    }

    private class CompletionOperation extends AsyncOperation {

        CompletionOperation(final AsyncOperation op) {
            super(op.stream, op.future, op.id);
        }

        @Override
        protected void doWork(final ServletOutputStream sos) {
            if (sos == null || sos.isReady()) {
                if (!future.isDone()) {
                    future.complete(null);
                }
            } else {
                // We need to requeue
                queueComplete(this);
            }
        }
    }

    protected final HttpServletResponse response;
    protected int status = 200;
    protected MultivaluedMap<String, Object> outputHeaders;
    protected final ResteasyProviderFactory factory;
    private OutputStream outputStream;
    protected volatile boolean suppressExceptionDuringChunkedTransfer = true;
    protected final HttpServletRequest request;
    protected final Map<Class<?>, Object> contextDataMap;

    // RESTEASY-1784
    @Override
    public void setSuppressExceptionDuringChunkedTransfer(boolean suppressExceptionDuringChunkedTransfer) {
        this.suppressExceptionDuringChunkedTransfer = suppressExceptionDuringChunkedTransfer;
    }

    @Override
    public boolean suppressExceptionDuringChunkedTransfer() {
        return suppressExceptionDuringChunkedTransfer;
    }

    /**
     * RESTEASY-684 wants to defer access to outputstream until a write happens
     *
     * <p>
     * Note that all locking is on {@code this} and should remain that way to avoid deadlocks on consumers of this
     * stream.
     * </p>
     *
     */
    protected class DeferredOutputStream extends AsyncOutputStream implements WriteListener {
        // Guarded by this
        private final Queue<AsyncOperation> asyncQueue;
        private final AtomicBoolean asyncRegistered;
        // Guarded by this
        private volatile boolean asyncListenerCalled;
        private volatile ServletOutputStream lazyOut;
        // Guarded by this
        private long idCounter;

        DeferredOutputStream() throws IOException {
            asyncQueue = new PriorityQueue<>(AsyncOperationComparator.INSTANCE);
            asyncRegistered = new AtomicBoolean();
        }

        @Override
        public void write(int i) throws IOException {
            getServletOutputStream().write(i);
        }

        @Override
        public void write(byte[] bytes) throws IOException {
            getServletOutputStream().write(bytes);
        }

        @Override
        public void write(byte[] bytes, int i, int i1) throws IOException {
            getServletOutputStream().write(bytes, i, i1);
        }

        @Override
        public void flush() throws IOException {
            getServletOutputStream().flush();
        }

        @Override
        public void close() throws IOException {
            //NOOP (RESTEASY-1650)
        }

        @Override
        public CompletionStage<Void> asyncFlush() {
            AsyncOperation op = new FlushOperation(this);
            queue(op);
            return op.future;
        }

        @Override
        public CompletionStage<Void> asyncWrite(byte[] bytes, int offset, int length) {
            AsyncOperation op = new WriteOperation(this, bytes, offset, length);
            queue(op);
            return op.future;
        }

        private void queue(AsyncOperation op) {
            // fetch it from the context directly to avoid having to restore the context just in case we're invoked on a context-less thread
            HttpRequest resteasyRequest = (HttpRequest) contextDataMap.get(HttpRequest.class);
            if (request.isAsyncStarted() && !resteasyRequest.getAsyncContext().isOnInitialRequest()) {
                boolean flush = false;
                final ServletOutputStream out;
                try {
                    out = getServletOutputStream();
                } catch (IOException e) {
                    op.future.completeExceptionally(e);
                    return;
                }
                if (asyncRegistered.compareAndSet(false, true)) {
                    out.setWriteListener(this);
                }
                synchronized (this) {
                    if (asyncListenerCalled && out.isReady()) {
                        // it's possible that we startAsync and queue, then queue another event and the stream becomes ready before
                        // onWritePossible is called, which means we need to flush the queue here to guarantee ordering if that happens
                        asyncQueue.add(op);
                        flush = true;
                    } else {
                        // just queue
                        asyncQueue.add(op);
                    }
                }
                // Invoked outside the lock to avoid deadlocks, the flushQueue itself locks on this
                if (flush) {
                    flushQueue();
                }
            } else {
                op.work(null);
            }
        }

        private void flushQueue() {
            synchronized (this) {
                final ServletOutputStream out;
                try {
                    out = getServletOutputStream();
                } catch (IOException e) {
                    onError(e);
                    return;
                }
                AsyncOperation op;
                while (out.isReady() && (op = asyncQueue.poll()) != null) {
                    op.work(out);
                }
            }
        }

        @Override
        public void onWritePossible() {
            asyncListenerCalled = true;
            flushQueue();
        }

        @Override
        public void onError(Throwable t) {
            synchronized (this) {
                asyncListenerCalled = true;
                AsyncOperation op;
                while ((op = asyncQueue.poll()) != null) {
                    if (!op.future.isDone())
                        op.future.completeExceptionally(t);
                }
            }
        }

        private ServletOutputStream getServletOutputStream() throws IOException {
            if (lazyOut == null) {
                synchronized (this) {
                    if (lazyOut == null) {
                        lazyOut = response.getOutputStream();
                    }
                }
            }
            return lazyOut;
        }

        private long getId() {
            synchronized (this) {
                if (idCounter == Long.MAX_VALUE) {
                    // This should never happen, but we will be safe and just reset the id in case it does.
                    idCounter = 0;
                }
                return idCounter++;
            }
        }
    }

    public HttpServletResponseWrapper(final HttpServletResponse response, final HttpServletRequest request,
            final ResteasyProviderFactory factory) {
        this.response = response;
        this.request = request;
        outputHeaders = new HttpServletResponseHeaders(response, factory);
        this.factory = factory;
        this.contextDataMap = ResteasyContext.getContextDataMap();
    }

    public int getStatus() {
        return status;
    }

    public void setStatus(int status) {
        this.status = status;
        this.response.setStatus(status);
    }

    public MultivaluedMap<String, Object> getOutputHeaders() {
        return outputHeaders;
    }

    public synchronized OutputStream getOutputStream() throws IOException {
        if (outputStream == null) {
            outputStream = new DeferredOutputStream();
        }
        return outputStream;
    }

    @Override
    public synchronized void setOutputStream(OutputStream os) {
        this.outputStream = os;
    }

    public void addNewCookie(NewCookie cookie) {
        outputHeaders.add(jakarta.ws.rs.core.HttpHeaders.SET_COOKIE, cookie);
    }

    public void sendError(int status) throws IOException {
        response.sendError(status);
    }

    public void sendError(int status, String message) throws IOException {
        response.sendError(status, message);
    }

    public boolean isCommitted() {
        return response.isCommitted();
    }

    public void reset() {
        response.reset();
        outputHeaders = new HttpServletResponseHeaders(response, factory);
    }

    @Override
    public void flushBuffer() throws IOException {
        response.flushBuffer();
    }

    private static class AsyncOperationComparator implements Comparator<AsyncOperation> {
        static final AsyncOperationComparator INSTANCE = new AsyncOperationComparator();

        @Override
        public int compare(final AsyncOperation o1, final AsyncOperation o2) {
            return Long.compare(o1.id, o2.id);
        }
    }
}
