/*
 * 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.websockets.spi;

import io.undertow.UndertowLogger;
import io.undertow.io.IoCallback;
import io.undertow.io.Sender;
import io.undertow.security.api.SecurityContext;
import io.undertow.security.idm.Account;
import io.undertow.server.HttpServerExchange;
import io.undertow.server.HttpUpgradeListener;
import io.undertow.server.session.SessionConfig;
import io.undertow.server.session.SessionManager;
import io.undertow.util.AttachmentKey;
import io.undertow.util.HeaderMap;
import io.undertow.util.HttpString;
import io.undertow.websockets.core.WebSocketChannel;
import org.xnio.ChannelListener;
import org.xnio.FinishedIoFuture;
import org.xnio.FutureResult;
import org.xnio.IoFuture;
import org.xnio.IoUtils;
import org.xnio.OptionMap;
import io.undertow.connector.ByteBufferPool;
import io.undertow.connector.PooledByteBuffer;
import org.xnio.channels.StreamSourceChannel;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.security.Principal;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Deque;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;

Author:
Stuart Douglas/**
 * @author Stuart Douglas
 */
public class AsyncWebSocketHttpServerExchange implements WebSocketHttpExchange {

    private final HttpServerExchange exchange;
    private Sender sender;
    private final Set<WebSocketChannel> peerConnections;

    public AsyncWebSocketHttpServerExchange(final HttpServerExchange exchange, Set<WebSocketChannel> peerConnections) {
        this.exchange = exchange;
        this.peerConnections = peerConnections;
    }


    @Override
    public <T> void putAttachment(final AttachmentKey<T> key, final T value) {
        exchange.putAttachment(key, value);
    }

    @Override
    public <T> T getAttachment(final AttachmentKey<T> key) {
        return exchange.getAttachment(key);
    }

    @Override
    public String getRequestHeader(final String headerName) {
        return exchange.getRequestHeaders().getFirst(HttpString.tryFromString(headerName));
    }

    @Override
    public Map<String, List<String>> getRequestHeaders() {
        Map<String, List<String>> headers = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
        for (final HttpString header : exchange.getRequestHeaders().getHeaderNames()) {
            headers.put(header.toString(), new ArrayList<>(exchange.getRequestHeaders().get(header)));
        }
        return Collections.unmodifiableMap(headers);
    }

    @Override
    public String getResponseHeader(final String headerName) {
        return exchange.getResponseHeaders().getFirst(HttpString.tryFromString(headerName));
    }

    @Override
    public Map<String, List<String>> getResponseHeaders() {
        Map<String, List<String>> headers = new HashMap<>();
        for (final HttpString header : exchange.getResponseHeaders().getHeaderNames()) {
            headers.put(header.toString(), new ArrayList<>(exchange.getResponseHeaders().get(header)));
        }
        return Collections.unmodifiableMap(headers);
    }

    @Override
    public void setResponseHeaders(final Map<String, List<String>> headers) {
        HeaderMap map = exchange.getRequestHeaders();
        map.clear();
        for (Map.Entry<String, List<String>> header : headers.entrySet()) {
            map.addAll(HttpString.tryFromString(header.getKey()), header.getValue());
        }
    }

    @Override
    public void setResponseHeader(final String headerName, final String headerValue) {
        exchange.getResponseHeaders().put(HttpString.tryFromString(headerName), headerValue);
    }

    @Override
    public void upgradeChannel(final HttpUpgradeListener upgradeCallback) {
        exchange.upgradeChannel(upgradeCallback);
    }

    @Override
    public IoFuture<Void> sendData(final ByteBuffer data) {
        if (sender == null) {
            this.sender = exchange.getResponseSender();
        }
        final FutureResult<Void> future = new FutureResult<>();
        sender.send(data, new IoCallback() {
            @Override
            public void onComplete(final HttpServerExchange exchange, final Sender sender) {
                future.setResult(null);
            }

            @Override
            public void onException(final HttpServerExchange exchange, final Sender sender, final IOException exception) {
                UndertowLogger.REQUEST_IO_LOGGER.ioException(exception);
                future.setException(exception);

            }
        });
        return future.getIoFuture();
    }

    @Override
    public IoFuture<byte[]> readRequestData() {
        final ByteArrayOutputStream data = new ByteArrayOutputStream();
        final PooledByteBuffer pooled = exchange.getConnection().getByteBufferPool().allocate();
        final ByteBuffer buffer = pooled.getBuffer();
        final StreamSourceChannel channel = exchange.getRequestChannel();
        int res;
        for (; ; ) {
            try {
                res = channel.read(buffer);
                if (res == -1) {
                    return new FinishedIoFuture<>(data.toByteArray());
                } else if (res == 0) {
                    //callback
                    final FutureResult<byte[]> future = new FutureResult<>();
                    channel.getReadSetter().set(new ChannelListener<StreamSourceChannel>() {
                        @Override
                        public void handleEvent(final StreamSourceChannel channel) {
                            int res;
                            try {
                                res = channel.read(buffer);
                                if (res == -1) {
                                    future.setResult(data.toByteArray());
                                    channel.suspendReads();
                                    return;
                                } else if (res == 0) {
                                    return;
                                } else {
                                    buffer.flip();
                                    while (buffer.hasRemaining()) {
                                        data.write(buffer.get());
                                    }
                                    buffer.clear();
                                }

                            } catch (IOException e) {
                                future.setException(e);
                            }
                        }
                    });
                    channel.resumeReads();
                    return future.getIoFuture();
                } else {
                    buffer.flip();
                    while (buffer.hasRemaining()) {
                        data.write(buffer.get());
                    }
                    buffer.clear();
                }

            } catch (IOException e) {
                final FutureResult<byte[]> future = new FutureResult<>();
                future.setException(e);
                return future.getIoFuture();
            }
        }


    }

    @Override
    public void endExchange() {
        exchange.endExchange();
    }

    @Override
    public void close() {
        try {
            exchange.endExchange();
        } finally {
            IoUtils.safeClose(exchange.getConnection());
        }
    }

    @Override
    public String getRequestScheme() {
        return exchange.getRequestScheme();
    }

    @Override
    public String getRequestURI() {
        String q = exchange.getQueryString();
        if (q == null || q.isEmpty()) {
            return exchange.getRequestURI();
        } else {
            return exchange.getRequestURI() + "?" + q;
        }
    }

    @Override
    public ByteBufferPool getBufferPool() {
        return exchange.getConnection().getByteBufferPool();
    }

    @Override
    public String getQueryString() {
        return exchange.getQueryString();
    }

    @Override
    public Object getSession() {
        SessionManager sm = exchange.getAttachment(SessionManager.ATTACHMENT_KEY);
        SessionConfig sessionCookieConfig = exchange.getAttachment(SessionConfig.ATTACHMENT_KEY);
        if(sm != null && sessionCookieConfig != null) {
            return sm.getSession(exchange, sessionCookieConfig);
        }
        return null;
    }

    @Override
    public Map<String, List<String>> getRequestParameters() {
        Map<String, List<String>> params = new HashMap<>();
        for (Map.Entry<String, Deque<String>> param : exchange.getQueryParameters().entrySet()) {
            params.put(param.getKey(), new ArrayList<>(param.getValue()));
        }
        return params;
    }

    @Override
    public Principal getUserPrincipal() {
        SecurityContext sc = exchange.getSecurityContext();
        if(sc == null) {
            return null;
        }
        Account authenticatedAccount = sc.getAuthenticatedAccount();
        if(authenticatedAccount == null) {
            return null;
        }
        return authenticatedAccount.getPrincipal();
    }

    @Override
    public boolean isUserInRole(String role) {
        SecurityContext sc = exchange.getSecurityContext();
        if(sc == null) {
            return false;
        }
        Account authenticatedAccount = sc.getAuthenticatedAccount();
        if(authenticatedAccount == null) {
            return false;
        }
        return authenticatedAccount.getRoles().contains(role);
    }

    @Override
    public Set<WebSocketChannel> getPeerConnections() {
        return peerConnections;
    }

    @Override
    public OptionMap getOptions() {
        return exchange.getConnection().getUndertowOptions();
    }
}