/*
 * Copyright 2015 LINE Corporation
 *
 * LINE Corporation 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 com.linecorp.armeria.client;

import static java.util.Objects.requireNonNull;

import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.util.Collections;
import java.util.EnumMap;
import java.util.Map;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;

import com.linecorp.armeria.client.pool.PoolKey;
import com.linecorp.armeria.common.SessionProtocol;
import com.linecorp.armeria.common.TimeoutException;

import io.netty.bootstrap.Bootstrap;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.EventLoop;
import io.netty.channel.pool.ChannelHealthChecker;
import io.netty.util.concurrent.Future;
import io.netty.util.concurrent.Promise;
import io.netty.util.internal.OneTimeTask;

class HttpSessionChannelFactory implements Function<PoolKey, Future<Channel>> {

    static final Object RETRY_WITH_H1C = new Object();

    static final ChannelHealthChecker HEALTH_CHECKER =
            ch -> ch.eventLoop().newSucceededFuture(HttpSessionHandler.get(ch).isActive());

    private final Bootstrap baseBootstrap;
    private final EventLoop eventLoop;
    private final Map<SessionProtocol, Bootstrap> bootstrapMap;
    private final RemoteInvokerOptions options;

    HttpSessionChannelFactory(Bootstrap bootstrap, RemoteInvokerOptions options) {
        baseBootstrap = requireNonNull(bootstrap);
        eventLoop = (EventLoop) bootstrap.group();

        bootstrapMap = Collections.synchronizedMap(new EnumMap<>(SessionProtocol.class));
        this.options = options;
    }

    @Override
    public Future<Channel> apply(PoolKey key) {
        final InetSocketAddress remoteAddress = key.remoteAddress();
        final SessionProtocol protocol = key.sessionProtocol();

        if (SessionProtocolNegotiationCache.isUnsupported(remoteAddress, protocol)) {
            // Fail immediately if it is sure that the remote address does not support the requested protocol.
            return eventLoop.newFailedFuture(
                    new SessionProtocolNegotiationException(protocol, "previously failed negotiation"));
        }

        final Promise<Channel> sessionPromise = eventLoop.newPromise();
        connect(remoteAddress, protocol, sessionPromise);

        return sessionPromise;
    }

    void connect(SocketAddress remoteAddress, SessionProtocol protocol, Promise<Channel> sessionPromise) {

        final Bootstrap bootstrap = bootstrap(protocol);
        final ChannelFuture connectFuture = bootstrap.connect(remoteAddress);
        final Channel ch = connectFuture.channel();

        if (connectFuture.isDone()) {
            notifySessionPromise(ch, connectFuture, sessionPromise);
        } else {
            connectFuture.addListener(
                    (Future<Void> future) -> notifySessionPromise(ch, future, sessionPromise));
        }
    }

    private Bootstrap bootstrap(SessionProtocol sessionProtocol) {
        return bootstrapMap.computeIfAbsent(sessionProtocol, sp -> {
            Bootstrap bs = baseBootstrap.clone();
            bs.handler(new ChannelInitializer<Channel>() {
                @Override
                protected void initChannel(Channel ch) throws Exception {
                    ch.pipeline().addLast(new HttpConfigurator(sp, options));
                }
            });
            return bs;
        });
    }

    private void notifySessionPromise(Channel ch, Future<Void> connectFuture,
                                      Promise<Channel> sessionPromise) {
        assert connectFuture.isDone();
        if (connectFuture.isSuccess()) {
            watchSessionActive(ch, sessionPromise);
        } else {
            sessionPromise.setFailure(connectFuture.cause());
        }
    }

    private Future<Channel> watchSessionActive(Channel ch, Promise<Channel> sessionPromise) {
        EventLoop eventLoop = ch.eventLoop();

        if (eventLoop.inEventLoop()) {
            watchSessionActive0(ch, sessionPromise);
        } else {
            eventLoop.execute(new OneTimeTask() {
                @Override
                public void run() {
                    watchSessionActive0(ch, sessionPromise);
                }
            });
        }
        return sessionPromise;
    }

    private void watchSessionActive0(final Channel ch, Promise<Channel> sessionPromise) {
        assert ch.eventLoop().inEventLoop();

        final ScheduledFuture<?> timeoutFuture = ch.eventLoop().schedule(new OneTimeTask() {
            @Override
            public void run() {
                if (sessionPromise.tryFailure(new TimeoutException(
                        "connection established, but session creation timed out: " + ch))) {
                    ch.close();
                }
            }
        }, options.connectTimeoutMillis(), TimeUnit.MILLISECONDS);

        ch.pipeline().addLast(new HttpSessionHandler(this, sessionPromise, timeoutFuture));
    }
}

@@ -1082,61 +1082,8 @@
 col;
-%0Aimport com.linecorp.armeria.common.TimeoutException;
 %0A%0Aim
@@ -3326,24 +3326,34 @@
 sionPromise(
+protocol, 
 ch, connectF
@@ -3368,32 +3368,32 @@
 essionPromise);%0A
-
         %7D else %7B
@@ -3498,16 +3498,26 @@
 Promise(
+protocol, 
 ch, futu
@@ -4096,46 +4096,44 @@
 ise(
-Channel ch, Future%3CVoid%3E connectFuture
+SessionProtocol protocol, Channel ch
 ,%0A  
@@ -4167,16 +4167,44 @@
         
+ Future%3CVoid%3E connectFuture,
  Promise
@@ -4342,16 +4342,26 @@
 nActive(
+protocol, 
 ch, sess
@@ -4516,27 +4516,100 @@
 nActive(
-Channel ch,
+SessionProtocol protocol, Channel ch,%0A                                              
  Promise
@@ -4628,32 +4628,33 @@
 ssionPromise) %7B%0A
+%0A
         EventLoo
@@ -4747,32 +4747,42 @@
 hSessionActive0(
+protocol, 
 ch, sessionPromi
@@ -4955,16 +4955,26 @@
 Active0(
+protocol, 
 ch, sess
@@ -5109,25 +5109,88 @@
 ve0(
-final Channel ch,
+SessionProtocol protocol, final Channel ch,%0A                                    
  Pro
@@ -5213,32 +5213,33 @@
 ssionPromise) %7B%0A
+%0A
         assert c
@@ -5465,23 +5465,42 @@
 ure(new 
-Timeout
+SessionProtocolNegotiation
 Exceptio
@@ -5525,16 +5525,26 @@
         
+ protocol,
  %22connec

Use SessionProtocolNegotiationException when protocol upgrade fails due to timeout