/*
 * Copyright The OpenZipkin Authors
 * SPDX-License-Identifier: Apache-2.0
 */
package brave.httpclient;

import brave.Span;
import brave.Tracer;
import brave.Tracer.SpanInScope;
import brave.http.HttpClientHandler;
import brave.http.HttpClientRequest;
import brave.http.HttpClientResponse;
import brave.http.HttpTracing;
import brave.internal.Nullable;
import brave.sampler.SamplerFunction;
import java.io.IOException;
import org.apache.http.Header;
import org.apache.http.HttpException;
import org.apache.http.HttpHost;
import org.apache.http.HttpRequest;
import org.apache.http.HttpResponse;
import org.apache.http.StatusLine;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpExecutionAware;
import org.apache.http.client.protocol.HttpClientContext;
import org.apache.http.conn.routing.HttpRoute;
import org.apache.http.impl.execchain.ClientExecChain;

import static brave.internal.Throwables.propagateIfFatal;

/**
 * protocol exec is the last in the execution chain, so first to execute. We eagerly create a span
 * here so that user interceptors can see it.
 */
final class TracingProtocolExec implements ClientExecChain {
  final Tracer tracer;
  final SamplerFunction<brave.http.HttpRequest> httpSampler;
  final HttpClientHandler<HttpClientRequest, HttpClientResponse> handler;
  final ClientExecChain protocolExec;

  TracingProtocolExec(HttpTracing httpTracing, ClientExecChain protocolExec) {
    this.tracer = httpTracing.tracing().tracer();
    this.httpSampler = httpTracing.clientRequestSampler();
    this.handler = HttpClientHandler.create(httpTracing);
    this.protocolExec = protocolExec;
  }

  @Override public CloseableHttpResponse execute(HttpRoute route,
    org.apache.http.client.methods.HttpRequestWrapper req,
    HttpClientContext context, HttpExecutionAware execAware)
    throws IOException, HttpException {
    HttpRequestWrapper request = new HttpRequestWrapper(req, context.getTargetHost());
    Span span = tracer.nextSpan(httpSampler, request);
    context.setAttribute(Span.class.getName(), span);

    CloseableHttpResponse response = null;
    Throwable error = null;
    SpanInScope scope = tracer.withSpanInScope(span);
    try {
      return response = protocolExec.execute(route, req, context, execAware);
    } catch (RuntimeException e) {
      error = e;
      throw e;
    } catch (HttpException e) {
      error = e;
      throw e;
    } catch (IOException e) {
      error = e;
      throw e;
    } catch (Error e) {
      propagateIfFatal(e);
      error = e;
      throw e;
    } finally {
      handler.handleReceive(new HttpResponseWrapper(response, context, error), span);
      scope.close();
    }
  }

  static final class HttpRequestWrapper extends HttpClientRequest {
    final HttpRequest request;
    @Nullable final HttpHost target;

    HttpRequestWrapper(HttpRequest request, @Nullable HttpHost target) {
      this.request = request;
      this.target = target;
    }

    @Override public Object unwrap() {
      return request;
    }

    @Override public String method() {
      return request.getRequestLine().getMethod();
    }

    @Override public String path() {
      if (request instanceof org.apache.http.client.methods.HttpRequestWrapper) {
        return ((org.apache.http.client.methods.HttpRequestWrapper) request).getURI().getPath();
      }
      String result = request.getRequestLine().getUri();
      int queryIndex = result.indexOf('?');
      return queryIndex == -1 ? result : result.substring(0, queryIndex);
    }

    @Override public String url() {
      if (target != null && request instanceof org.apache.http.client.methods.HttpRequestWrapper) {
        org.apache.http.client.methods.HttpRequestWrapper
          wrapper = (org.apache.http.client.methods.HttpRequestWrapper) request;
        return target.toURI() + wrapper.getURI();
      }
      return request.getRequestLine().getUri();
    }

    @Override public String header(String name) {
      Header result = request.getFirstHeader(name);
      return result != null ? result.getValue() : null;
    }

    @Override public void header(String name, String value) {
      request.setHeader(name, value);
    }
  }

  static final class HttpResponseWrapper extends HttpClientResponse {
    @Nullable final HttpRequestWrapper request;
    @Nullable final HttpResponse response;
    @Nullable final Throwable error;

    HttpResponseWrapper(@Nullable HttpResponse response, HttpClientContext context,
      @Nullable Throwable error) {
      HttpRequest request = context.getRequest();
      HttpHost target = context.getTargetHost();
      this.request = request != null ? new HttpRequestWrapper(request, target) : null;
      this.response = response;
      this.error = error;
    }

    @Override public Object unwrap() {
      return response;
    }

    @Override @Nullable public HttpRequestWrapper request() {
      return request;
    }

    @Override public Throwable error() {
      return error;
    }

    @Override public int statusCode() {
      if (response == null) return 0;
      StatusLine statusLine = response.getStatusLine();
      return statusLine != null ? statusLine.getStatusCode() : 0;
    }
  }
}
