/*
 * Copyright The OpenTelemetry Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package io.opentelemetry.instrumentation.ktor.v2_0.client

import io.ktor.client.*
import io.ktor.client.call.*
import io.ktor.client.plugins.*
import io.ktor.client.request.*
import io.ktor.client.statement.*
import io.ktor.util.*
import io.ktor.util.pipeline.*
import io.opentelemetry.context.Context
import io.opentelemetry.context.propagation.ContextPropagators
import io.opentelemetry.extension.kotlin.asContextElement
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter
import kotlinx.coroutines.withContext

class KtorClientTracing internal constructor(
  private val instrumenter: Instrumenter<HttpRequestData, HttpResponse>,
  private val propagators: ContextPropagators,
) {

  private fun createSpan(requestBuilder: HttpRequestBuilder): Context? {
    val parentContext = Context.current()
    val requestData = requestBuilder.build()

    return if (instrumenter.shouldStart(parentContext, requestData)) {
      instrumenter.start(parentContext, requestData)
    } else {
      null
    }
  }

  private fun populateRequestHeaders(requestBuilder: HttpRequestBuilder, context: Context) {
    propagators.textMapPropagator.inject(context, requestBuilder, KtorHttpHeadersSetter)
  }

  private fun endSpan(context: Context, call: HttpClientCall, error: Throwable?) {
    endSpan(context, HttpRequestBuilder().takeFrom(call.request), call.response, error)
  }

  private fun endSpan(context: Context, requestBuilder: HttpRequestBuilder, response: HttpResponse?, error: Throwable?) {
    instrumenter.end(context, requestBuilder.build(), response, error)
  }

  companion object : HttpClientPlugin<KtorClientTracingBuilder, KtorClientTracing> {

    private val openTelemetryContextKey = AttributeKey<Context>("OpenTelemetry")

    override val key = AttributeKey<KtorClientTracing>("OpenTelemetry")

    override fun prepare(block: KtorClientTracingBuilder.() -> Unit) = KtorClientTracingBuilder().apply(block).build()

    override fun install(plugin: KtorClientTracing, scope: HttpClient) {
      installSpanCreation(plugin, scope)
      installSpanEnd(plugin, scope)
    }

    private fun installSpanCreation(plugin: KtorClientTracing, scope: HttpClient) {
      val createSpanPhase = PipelinePhase("OpenTelemetryCreateSpan")
      scope.sendPipeline.insertPhaseAfter(HttpSendPipeline.State, createSpanPhase)

      scope.sendPipeline.intercept(createSpanPhase) {
        val requestBuilder = context
        val openTelemetryContext = plugin.createSpan(requestBuilder)

        if (openTelemetryContext != null) {
          try {
            requestBuilder.attributes.put(openTelemetryContextKey, openTelemetryContext)
            plugin.populateRequestHeaders(requestBuilder, openTelemetryContext)

            withContext(openTelemetryContext.asContextElement()) { proceed() }
          } catch (e: Throwable) {
            plugin.endSpan(openTelemetryContext, requestBuilder, null, e)
            throw e
          }
        } else {
          proceed()
        }
      }
    }

    private fun installSpanEnd(plugin: KtorClientTracing, scope: HttpClient) {
      val endSpanPhase = PipelinePhase("OpenTelemetryEndSpan")
      scope.receivePipeline.insertPhaseBefore(HttpReceivePipeline.State, endSpanPhase)

      scope.receivePipeline.intercept(endSpanPhase) {
        val openTelemetryContext = it.call.attributes.getOrNull(openTelemetryContextKey)

        if (openTelemetryContext != null) {
          try {
            withContext(openTelemetryContext.asContextElement()) { proceed() }
            plugin.endSpan(openTelemetryContext, it.call, null)
          } catch (e: Throwable) {
            plugin.endSpan(openTelemetryContext, it.call, e)
            throw e
          }
        } else {
          proceed()
        }
      }
    }
  }
}
