package org.chenyang.http.impl.client;

import org.chenyang.http.ConnectionReuseStrategy;
import org.chenyang.http.HttpEntity;
import org.chenyang.http.HttpException;
import org.chenyang.http.HttpHost;
import org.chenyang.http.HttpRequest;
import org.chenyang.http.HttpRequestInterceptor;
import org.chenyang.http.HttpResponse;
import org.chenyang.http.HttpVersion;
import org.chenyang.http.auth.AuthSchemeRegistry;
import org.chenyang.http.auth.AuthScope;
import org.chenyang.http.auth.AuthState;
import org.chenyang.http.auth.Credentials;
import org.chenyang.http.client.config.RequestConfig;
import org.chenyang.http.client.params.HttpClientParamConfig;
import org.chenyang.http.client.protocol.RequestClientConnControl;
import org.chenyang.http.config.ConnectionConfig;
import org.chenyang.http.conn.HttpConnectionFactory;
import org.chenyang.http.conn.ManagedHttpClientConnection;
import org.chenyang.http.conn.routing.HttpRoute;
import org.chenyang.http.conn.routing.RouteInfo;
import org.chenyang.http.entity.BufferedHttpEntity;
import org.chenyang.http.impl.DefaultConnectionReuseStrategy;
import org.chenyang.http.impl.auth.BasicSchemeFactory;
import org.chenyang.http.impl.auth.DigestSchemeFactory;
import org.chenyang.http.impl.auth.KerberosSchemeFactory;
import org.chenyang.http.impl.auth.NTLMSchemeFactory;
import org.chenyang.http.impl.auth.SPNegoSchemeFactory;
import org.chenyang.http.impl.conn.ManagedHttpClientConnectionFactory;
import org.chenyang.http.message.BasicHttpRequest;
import org.chenyang.http.params.BasicHttpParams;
import org.chenyang.http.params.HttpParamConfig;
import org.chenyang.http.params.HttpParams;
import org.chenyang.http.protocol.BasicHttpContext;
import org.chenyang.http.protocol.HttpContext;
import org.chenyang.http.protocol.HttpProcessor;
import org.chenyang.http.protocol.HttpRequestExecutor;
import org.chenyang.http.protocol.ImmutableHttpProcessor;
import org.chenyang.http.protocol.RequestTargetHost;
import org.chenyang.http.protocol.RequestUserAgent;
import org.chenyang.http.util.Args;
import org.chenyang.http.util.EntityUtils;
import org.chenyang.http.impl.auth.HttpAuthenticator;
import org.chenyang.http.impl.execchain.TunnelRefusedException;

import java.io.IOException;
import java.net.Socket;

public class ProxyClient {
   private final HttpConnectionFactory<HttpRoute, ManagedHttpClientConnection> connFactory;
   private final ConnectionConfig connectionConfig;
   private final RequestConfig requestConfig;
   private final HttpProcessor httpProcessor;
   private final HttpRequestExecutor requestExec;
   private final ProxyAuthenticationStrategy proxyAuthStrategy;
   private final HttpAuthenticator authenticator;
   private final AuthState proxyAuthState;
   private final AuthSchemeRegistry authSchemeRegistry;
   private final ConnectionReuseStrategy reuseStrategy;

   public ProxyClient(HttpConnectionFactory<HttpRoute, ManagedHttpClientConnection> connFactory, ConnectionConfig connectionConfig, RequestConfig requestConfig) {
      this.connFactory = (HttpConnectionFactory<HttpRoute, ManagedHttpClientConnection>)(connFactory != null ? connFactory : ManagedHttpClientConnectionFactory.INSTANCE);
      this.connectionConfig = connectionConfig != null ? connectionConfig : ConnectionConfig.DEFAULT;
      this.requestConfig = requestConfig != null ? requestConfig : RequestConfig.DEFAULT;
      this.httpProcessor = new ImmutableHttpProcessor(new HttpRequestInterceptor[]{new RequestTargetHost(), new RequestClientConnControl(), new RequestUserAgent()});
      this.requestExec = new HttpRequestExecutor();
      this.proxyAuthStrategy = new ProxyAuthenticationStrategy();
      this.authenticator = new HttpAuthenticator();
      this.proxyAuthState = new AuthState();
      this.authSchemeRegistry = new AuthSchemeRegistry();
      this.authSchemeRegistry.register("Basic", new BasicSchemeFactory());
      this.authSchemeRegistry.register("Digest", new DigestSchemeFactory());
      this.authSchemeRegistry.register("NTLM", new NTLMSchemeFactory());
      this.authSchemeRegistry.register("Negotiate", new SPNegoSchemeFactory());
      this.authSchemeRegistry.register("Kerberos", new KerberosSchemeFactory());
      this.reuseStrategy = new DefaultConnectionReuseStrategy();
   }

   /** @deprecated */
   @Deprecated
   public ProxyClient(HttpParams params) {
      this((HttpConnectionFactory)null, HttpParamConfig.getConnectionConfig(params), HttpClientParamConfig.getRequestConfig(params));
   }

   public ProxyClient(RequestConfig requestConfig) {
      this((HttpConnectionFactory)null, (ConnectionConfig)null, requestConfig);
   }

   public ProxyClient() {
      this((HttpConnectionFactory)null, (ConnectionConfig)null, (RequestConfig)null);
   }

   /** @deprecated */
   @Deprecated
   public HttpParams getParams() {
      return new BasicHttpParams();
   }

   /** @deprecated */
   @Deprecated
   public AuthSchemeRegistry getAuthSchemeRegistry() {
      return this.authSchemeRegistry;
   }

   public Socket tunnel(HttpHost proxy, HttpHost target, Credentials credentials) throws IOException, HttpException {
      Args.notNull(proxy, "Proxy host");
      Args.notNull(target, "Target host");
      Args.notNull(credentials, "Credentials");
      HttpHost host = target;
      if (target.getPort() <= 0) {
         host = new HttpHost(target.getHostName(), 80, target.getSchemeName());
      }

      HttpRoute route = new HttpRoute(host, this.requestConfig.getLocalAddress(), proxy, false, RouteInfo.TunnelType.TUNNELLED, RouteInfo.LayerType.PLAIN);
      ManagedHttpClientConnection conn = this.connFactory.create(route, this.connectionConfig);
      HttpContext context = new BasicHttpContext();
      HttpRequest connect = new BasicHttpRequest("CONNECT", host.toHostString(), HttpVersion.HTTP_1_1);
      BasicCredentialsProvider credsProvider = new BasicCredentialsProvider();
      credsProvider.setCredentials(new AuthScope(proxy), credentials);
      context.setAttribute("http.target_host", target);
      context.setAttribute("http.connection", conn);
      context.setAttribute("http.request", connect);
      context.setAttribute("http.route", route);
      context.setAttribute("http.auth.proxy-scope", this.proxyAuthState);
      context.setAttribute("http.auth.credentials-provider", credsProvider);
      context.setAttribute("http.authscheme-registry", this.authSchemeRegistry);
      context.setAttribute("http.request-config", this.requestConfig);
      this.requestExec.preProcess(connect, this.httpProcessor, context);

      while(true) {
         if (!conn.isOpen()) {
            Socket socket = new Socket(proxy.getHostName(), proxy.getPort());
            conn.bind(socket);
         }

         this.authenticator.generateAuthResponse(connect, this.proxyAuthState, context);
         HttpResponse response = this.requestExec.execute(connect, conn, context);
         int status = response.getStatusLine().getStatusCode();
         if (status < 200) {
            throw new HttpException("Unexpected response to CONNECT request: " + response.getStatusLine());
         }

         if (!this.authenticator.isAuthenticationRequested(proxy, response, this.proxyAuthStrategy, this.proxyAuthState, context) || !this.authenticator.handleAuthChallenge(proxy, response, this.proxyAuthStrategy, this.proxyAuthState, context)) {
            status = response.getStatusLine().getStatusCode();
            if (status > 299) {
               HttpEntity entity = response.getEntity();
               if (entity != null) {
                  response.setEntity(new BufferedHttpEntity(entity));
               }

               conn.close();
               throw new TunnelRefusedException("CONNECT refused by proxy: " + response.getStatusLine(), response);
            } else {
               return conn.getSocket();
            }
         }

         if (this.reuseStrategy.keepAlive(response, context)) {
            HttpEntity entity = response.getEntity();
            EntityUtils.consume(entity);
         } else {
            conn.close();
         }

         connect.removeHeaders("Proxy-Authorization");
      }
   }
}
