/*
  Copyright (C) 2013-2018 Expedia Inc.

  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 com.hotels.styx.client;

import com.hotels.styx.api.LiveHttpRequest;
import com.hotels.styx.api.LiveHttpResponse;
import com.hotels.styx.api.Id;
import com.hotels.styx.api.ResponseEventListener;
import com.hotels.styx.api.exceptions.NoAvailableHostsException;
import com.hotels.styx.client.connectionpool.ConnectionPool;
import rx.Observable;

import java.util.Optional;
import java.util.concurrent.atomic.AtomicReference;

import static java.util.Objects.requireNonNull;

/**
 * Encapsulates a single connection to remote server which we can use to send the messages.
 */
class Transport {
    private final Id appId;
    private final CharSequence originIdHeaderName;

    public Transport(Id appId, CharSequence originIdHeaderName) {
        this.appId = requireNonNull(appId);
        this.originIdHeaderName = requireNonNull(originIdHeaderName);
    }

    public HttpTransaction send(LiveHttpRequest request, Optional<ConnectionPool> origin, Id originId) {
        Observable<Connection> connection = connection(request, origin);

        AtomicReference<Connection> connectionRef = new AtomicReference<>(null);
        Observable<LiveHttpResponse> observableResponse = connection.flatMap(tConnection -> {
            connectionRef.set(tConnection);
            return tConnection.write(request)
                    .map(response -> addOriginId(originId, response));
        });

        return new HttpTransaction() {
            @Override
            public Observable<LiveHttpResponse> response() {
                return ResponseEventListener.from(observableResponse)
                        .whenCancelled(() -> closeIfConnected(origin, connectionRef))
                        .whenResponseError(cause -> closeIfConnected(origin, connectionRef))
                        .whenContentError(cause -> closeIfConnected(origin, connectionRef))
                        .whenCompleted(() -> returnIfConnected(origin, connectionRef))
                        .apply();
            }

            private synchronized void closeIfConnected(Optional<ConnectionPool> connectionPool, AtomicReference<Connection> connectionRef) {
                Connection connection = connectionRef.get();
                if (connection != null && connectionPool.isPresent()) {
                    connectionPool.get().closeConnection(connection);
                    connectionRef.set(null);
                }
            }

            private synchronized void returnIfConnected(Optional<ConnectionPool> connectionPool, AtomicReference<Connection> connectionRef) {
                Connection connection = connectionRef.get();
                if (connection != null && connectionPool.isPresent()) {
                    connectionPool.get().returnConnection(connection);
                    connectionRef.set(null);
                }
            }
        };
    }

    private Observable<Connection> connection(LiveHttpRequest request, Optional<ConnectionPool> origin) {
        return origin
                .map(ConnectionPool::borrowConnection)
                .orElseGet(() -> {
                    // Aggregates an empty body:
                    request.consume();
                    return Observable.error(new NoAvailableHostsException(appId));
                });
    }

    private LiveHttpResponse addOriginId(Id originId, LiveHttpResponse response) {
        return response.newBuilder()
                .header(originIdHeaderName, originId)
                .build();
    }
}

@@ -611,16 +611,47 @@
 lient;%0A%0A
+import com.hotels.styx.api.Id;%0A
 import c
@@ -735,39 +735,8 @@
 se;%0A
-import com.hotels.styx.api.Id;%0A
 impo
@@ -960,60 +960,8 @@
 nal;
-%0Aimport java.util.concurrent.atomic.AtomicReference;
 %0A%0Aim
@@ -1583,88 +1583,79 @@
     
-AtomicReference%3CConnection%3E connectionRef = new AtomicReference%3C%3E(null);%0A
+return new HttpTransaction() %7B%0A            @Override%0A     
        
+public
  Obs
@@ -1684,28 +1684,43 @@
 se%3E 
-observableResponse =
+response() %7B%0A                return
  con
@@ -1735,18 +1735,17 @@
 flatMap(
-tC
+c
 onnectio
@@ -1767,61 +1767,67 @@
     
-connectionRef.set(tConnection);%0A            return tC
+        Observable%3CLiveHttpResponse%3E responseObservable = c
 onne
@@ -1847,16 +1847,24 @@
 equest)%0A
+        
         
@@ -1930,143 +1930,13 @@
 ));%0A
+%0A
     
-    %7D);%0A%0A        return new HttpTransaction() %7B%0A            @Override%0A            public Observable%3CLiveHttpResponse%3E response() %7B%0A
     
@@ -1985,28 +1985,32 @@
 rom(
-observableResponse)%0A
+responseObservable)%0A    
     
@@ -2077,35 +2077,32 @@
 igin, connection
-Ref
 ))%0A             
@@ -2108,24 +2108,28 @@
             
+    
 .whenRespons
@@ -2171,38 +2171,39 @@
 igin, connection
-Ref
 ))%0A
+    
                 
@@ -2264,38 +2264,39 @@
 igin, connection
-Ref
 ))%0A
+    
                 
@@ -2360,19 +2360,16 @@
 nnection
-Ref
 ))%0A     
@@ -2387,16 +2387,20 @@
         
+    
 .apply()
@@ -2397,24 +2397,44 @@
   .apply();%0A
+                %7D);%0A
             
@@ -2532,32 +2532,16 @@
 onPool, 
-AtomicReference%3C
 Connecti
@@ -2534,33 +2534,32 @@
 Pool, Connection
-%3E
  connectionRef) 
@@ -2557,76 +2557,12 @@
 tion
-Ref
 ) %7B%0A
-                Connection connection = connectionRef.get();%0A
     
@@ -2703,53 +2703,8 @@
 n);%0A
-                    connectionRef.set(null);%0A
     
@@ -2833,24 +2833,8 @@
 ol, 
-AtomicReference%3C
 Conn
@@ -2831,33 +2831,32 @@
 Pool, Connection
-%3E
  connectionRef) 
@@ -2854,76 +2854,12 @@
 tion
-Ref
 ) %7B%0A
-                Connection connection = connectionRef.get();%0A
     
@@ -3001,53 +3001,8 @@
 n);%0A
-                    connectionRef.set(null);%0A
     

Remove mutable state (connectionRef) from Transport class. (#329)