const {Request, Annotation} = require('zipkin');
const url = require('url');

class ExpressHttpProxyInstrumentation {

  constructor({tracer, serviceName = tracer.localEndpoint.serviceName, remoteServiceName}) {
    this.tracer = tracer;
    this.serviceName = serviceName;
    this.remoteServiceName = remoteServiceName;
  }

  decorateAndRecordRequest(serverReq, proxyReq, serverTraceId) {

    return this.tracer.letId(serverTraceId, () => {
      const clientTraceId = this.tracer.createChildId();
      this.tracer.setId(clientTraceId);

      const proxyReqWithZipkinHeaders = Request.addZipkinHeaders(proxyReq, clientTraceId);
      Object.defineProperty(
        serverReq, 
        '_trace_id_proxy',
        {
          configurable: false, 
          get: () => clientTraceId
        }
      );
      this._recordRequest(proxyReqWithZipkinHeaders);

      return proxyReqWithZipkinHeaders;
    });
  }

  _recordRequest(proxyReq) {
    this.tracer.recordServiceName(this.serviceName);
    this.tracer.recordRpc(proxyReq.method.toUpperCase());
    this.tracer.recordBinary('http.path', url.parse(proxyReq.path).pathname);
    this.tracer.recordAnnotation(new Annotation.ClientSend());
    if (this.remoteServiceName) {
      this.tracer.recordAnnotation(new Annotation.ServerAddr({
        serviceName: this.remoteServiceName,
        port: parseInt(proxyReq.port)
      }));
    }
  }

  recordResponse(rsp, clientTraceId) {
    this.tracer.letId(clientTraceId, () => {
      this.tracer.recordBinary('http.status_code', rsp.statusCode.toString());
      this.tracer.recordAnnotation(new Annotation.ClientRecv());
    });
  }
}

const wrapProxyReqOptDecorator = (instrumentation, proxyReqOptDecorator) => {

  return (proxyReqOpts, serverReq) => {
    const serverTraceId = serverReq._trace_id;
    let wrappedProxyReqOpts = proxyReqOpts;

    if (typeof proxyReqOptDecorator === 'function') {
      tracer.letId(serverTraceId, () => {
          wrappedProxyReqOpts = proxyReqOptDecorator(proxyReqOpts, serverReq);
      });
    }

    return instrumentation.decorateAndRecordRequest(serverReq, wrappedProxyReqOpts, serverTraceId);
  };
}

const wrapUserResHeaderDecorator = (instrumentation, userResHeaderDecorator) => {

  return (headers, userReq, userRes, proxyReq, proxyRes) => {
    const serverTraceId = userReq._trace_id;
    let wrappedHeaders = headers;

    if (typeof userResHeaderDecorator === 'function') {
      tracer.letId(serverTraceId, () => {
        wrappedHeaders = userResHeaderDecorator(headers, userReq, userRes, proxyReq, proxyRes)
      }); 
    }
    instrumentation.recordResponse(proxyRes, userReq._trace_id_proxy);

    return wrappedHeaders;
  };
}

const wrapProxy = (proxy, {tracer, serviceName, remoteServiceName}) => {

  return function zipkinProxy(host, options = {}) {

    const instrumentation = new ExpressHttpProxyInstrumentation({tracer, serviceName, remoteServiceName});
    const wrappedOptions = options;

    const {proxyReqOptDecorator} = wrappedOptions;
    wrappedOptions.proxyReqOptDecorator = wrapProxyReqOptDecorator(instrumentation, proxyReqOptDecorator);

    const {userResHeaderDecorator} = wrappedOptions;
    wrappedOptions.userResHeaderDecorator = wrapUserResHeaderDecorator(instrumentation, userResHeaderDecorator);

    return proxy(host, wrappedOptions);
  };
}

module.exports = wrapProxy;