package unitel

import (
	"context"
	"net/http"

	"github.com/getsentry/sentry-go"
	"github.com/rs/zerolog/log"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/trace"
)

type contextKey struct {
	Key string
}

const (
	SentryTraceHeader   = "sentry-trace"
	SentryBaggageHeader = "sentry-baggage"
)

var (
	tracerContextKey = contextKey{"tracer"}
)

type SpanStartOpt = func(context.Context) (context.Context, []trace.SpanStartOption, []sentry.SpanOption)

func (t *Telemetry) StartSpan(ctx context.Context, operation, name string, cfgs ...SpanStartOpt) *Span {
	otelStartOpts := make([]trace.SpanStartOption, 0)
	sentryStartOpts := []sentry.SpanOption{sentry.WithTransactionName(name), sentry.WithDescription(name)}

	for _, opt := range cfgs {
		var otel []trace.SpanStartOption
		var sentry []sentry.SpanOption

		ctx, otel, sentry = opt(ctx)

		otelStartOpts = append(otelStartOpts, otel...)
		sentryStartOpts = append(sentryStartOpts, sentry...)
	}

	var otelTracer trace.Tracer
	if tracer, ok := ctx.Value(tracerContextKey).(trace.Tracer); ok && tracer != nil {
		otelTracer = tracer
	} else {
		otelTracer = t.Tracer
	}

	ctx, otelSpan := otelTracer.Start(ctx, name, otelStartOpts...)
	sentrySpan := sentry.StartSpan(ctx, operation, sentryStartOpts...)

	return &Span{
		otel:   otelSpan,
		sentry: sentrySpan,
	}
}

func WithOtelOptions(opts ...trace.SpanStartOption) SpanStartOpt {
	return func(ctx context.Context) (context.Context, []trace.SpanStartOption, []sentry.SpanOption) {
		return ctx, opts, []sentry.SpanOption{}
	}
}

func (t *Telemetry) InjectIntoHeaders(ctx context.Context, headers http.Header) {
	t.Propagator.Inject(ctx, propagation.HeaderCarrier(headers))

	if sentrySpan := sentry.SpanFromContext(ctx); sentrySpan == nil {
		// this should never happen...
		log.Trace().Msgf("failed to inject Sentry span ID, Sentry span could not be extracted from the context.Context")
	} else {
		headers.Set(SentryTraceHeader, sentrySpan.ToSentryTrace())
		headers.Set(SentryBaggageHeader, sentrySpan.ToBaggage())
	}
}

func (t *Telemetry) InjectIntoMap(ctx context.Context, m map[string]string) {
	t.Propagator.Inject(ctx, propagation.MapCarrier(m))

	if sentrySpan := sentry.SpanFromContext(ctx); sentrySpan == nil {
		// this should never happen...
		log.Trace().Msgf("failed to inject Sentry span ID, Sentry span could not be extracted from the context.Context")
	} else {
		m[SentryTraceHeader] = sentrySpan.ToSentryTrace()
		m[SentryBaggageHeader] = sentrySpan.ToBaggage()
	}
}

func (t *Telemetry) ContinueFromHeaders(h http.Header) SpanStartOpt {
	return func(ctx context.Context) (context.Context, []trace.SpanStartOption, []sentry.SpanOption) {
		ctx = t.Propagator.Extract(ctx, propagation.HeaderCarrier(h))

		return ctx,
			[]trace.SpanStartOption{},
			[]sentry.SpanOption{
				sentry.ContinueFromHeaders(h.Get(sentry.SentryTraceHeader), h.Get(sentry.SentryBaggageHeader)),
				sentry.WithTransactionSource(sentry.SourceURL),
			}
	}
}

func (t *Telemetry) ContinueFromMap(m map[string]string) SpanStartOpt {
	return func(ctx context.Context) (context.Context, []trace.SpanStartOption, []sentry.SpanOption) {
		ctx = t.Propagator.Extract(ctx, propagation.MapCarrier(m))

		sentryTrace := m[sentry.SentryTraceHeader]
		sentryBaggage := m[sentry.SentryBaggageHeader]

		return ctx,
			[]trace.SpanStartOption{},
			[]sentry.SpanOption{
				sentry.ContinueFromHeaders(sentryTrace, sentryBaggage),
				sentry.WithTransactionSource(sentry.SourceURL),
			}
	}
}

func WithOtelTracer(tracer trace.Tracer) SpanStartOpt {
	return func(ctx context.Context) (context.Context, []trace.SpanStartOption, []sentry.SpanOption) {
		return context.WithValue(ctx, tracerContextKey, tracer), []trace.SpanStartOption{}, []sentry.SpanOption{}
	}
}

type TracedFunc[T any] func(context.Context, string, ...any) TracedFuncResult[T]
type TracedFuncResult[T any] struct {
	Result T
	Err    error
}

func RunTraced[T any](t *Telemetry, op string, fn func(context.Context, ...any) TracedFuncResult[T]) TracedFunc[T] {
	return func(ctx context.Context, name string, args ...any) TracedFuncResult[T] {
		span := t.StartSpan(ctx, op, name)
		defer span.End()
		ctx = span.Context()

		return fn(ctx, args...)
	}
}