/*
 * JBoss, Home of Professional Open Source.
 * Copyright 2014 Red Hat, Inc., and individual contributors
 * as indicated by the @author tags.
 *
 * Licensed 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 io.undertow.server.handlers;

import io.undertow.UndertowLogger;
import io.undertow.connector.PooledByteBuffer;
import io.undertow.server.Connectors;
import io.undertow.server.HandlerWrapper;
import io.undertow.server.HttpHandler;
import io.undertow.server.HttpServerExchange;
import io.undertow.server.handlers.builder.HandlerBuilder;
import io.undertow.server.protocol.http.HttpContinue;
import org.xnio.ChannelListener;
import org.xnio.IoUtils;
import org.xnio.channels.StreamSourceChannel;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.Map;
import java.util.Set;
/**
 *
 * Handler that will buffer all request data
 *
 * @author Stuart Douglas
 */
public class RequestBufferingHandler implements HttpHandler {

    private final HttpHandler next;
    private final int maxBuffers;

    public RequestBufferingHandler(HttpHandler next, int maxBuffers) {
        this.next = next;
        this.maxBuffers = maxBuffers;
    }

    @Override
    public void handleRequest(final HttpServerExchange exchange) throws Exception {

        if(!exchange.isRequestComplete() && !HttpContinue.requiresContinueResponse(exchange.getRequestHeaders())) {
            final StreamSourceChannel channel = exchange.getRequestChannel();
            int readBuffers = 0;
            final PooledByteBuffer[] bufferedData = new PooledByteBuffer[maxBuffers];
            PooledByteBuffer buffer = exchange.getConnection().getByteBufferPool().allocate();
            do {
                int r;
                ByteBuffer b = buffer.getBuffer();
                r = channel.read(b);
                if (r == -1) { //TODO: listener read
                    if (b.position() == 0) {
                        buffer.close();
                    } else {
                        b.flip();
                        bufferedData[readBuffers] = buffer;
                    }
                    break;
                } else if(r == 0) {
                    final PooledByteBuffer finalBuffer = buffer;
                    final int finalReadBuffers = readBuffers;
                    channel.getReadSetter().set(new ChannelListener<StreamSourceChannel>() {

                        PooledByteBuffer buffer = finalBuffer;
                        int readBuffers = finalReadBuffers;
                        @Override
                        public void handleEvent(StreamSourceChannel channel) {
                            try {
                                do {
                                    int r;
                                    ByteBuffer b = buffer.getBuffer();
                                    r = channel.read(b);
                                    if (r == -1) { //TODO: listener read
                                        if (b.position() == 0) {
                                            buffer.close();
                                        } else {
                                            b.flip();
                                            bufferedData[readBuffers] = buffer;
                                        }
                                        Connectors.ungetRequestBytes(exchange, bufferedData);
                                        Connectors.resetRequestChannel(exchange);
                                        Connectors.executeRootHandler(next, exchange);
                                        channel.getReadSetter().set(null);
                                        return;
                                    } else if (r == 0) {
                                        return;
                                    } else if (!b.hasRemaining()) {
                                        b.flip();
                                        bufferedData[readBuffers++] = buffer;
                                        if (readBuffers == maxBuffers) {
                                            Connectors.ungetRequestBytes(exchange, bufferedData);
                                            Connectors.resetRequestChannel(exchange);
                                            Connectors.executeRootHandler(next, exchange);
                                            channel.getReadSetter().set(null);
                                            return;
                                        }
                                        buffer = exchange.getConnection().getByteBufferPool().allocate();
                                    }
                                } while (true);
                            } catch (IOException e) {
                                UndertowLogger.REQUEST_IO_LOGGER.ioException(e);
                                for(int i = 0; i < bufferedData.length; ++i) {
                                    IoUtils.safeClose(bufferedData[i]);
                                }
                                exchange.endExchange();
                            }
                        }
                    });
                    channel.resumeReads();
                    return;
                } else if (!b.hasRemaining()) {
                    b.flip();
                    bufferedData[readBuffers++] = buffer;
                    if(readBuffers == maxBuffers) {
                        break;
                    }
                    buffer = exchange.getConnection().getByteBufferPool().allocate();
                }
            } while (true);
            Connectors.ungetRequestBytes(exchange, bufferedData);
            Connectors.resetRequestChannel(exchange);
        }
        next.handleRequest(exchange);
    }

    public static final class Wrapper implements HandlerWrapper {

        private final int maxBuffers;

        public Wrapper(int maxBuffers) {
            this.maxBuffers = maxBuffers;
        }

        @Override
        public HttpHandler wrap(HttpHandler handler) {
            return new RequestBufferingHandler(handler, maxBuffers);
        }
    }

    public static final class Builder implements HandlerBuilder {

        @Override
        public String name() {
            return "buffer-request";
        }

        @Override
        public Map<String, Class<?>> parameters() {
            return Collections.<String, Class<?>>singletonMap("buffers", Integer.class);
        }

        @Override
        public Set<String> requiredParameters() {
            return Collections.singleton("buffers");
        }

        @Override
        public String defaultParameter() {
            return "buffers";
        }

        @Override
        public HandlerWrapper build(Map<String, Object> config) {
            return new Wrapper((Integer) config.get("buffers"));
        }
    }
}
