package fiber_tracing

// This was copied from "github.com/psmarcin/fiber-opentelemetry/pkg/fiber-otel"
// and slighltly modified but this piece of code is yet to be fully integrated
// into the package.

import (
	"git.ma-al.com/maal-libraries/observer/pkg/attr"
	"github.com/gofiber/fiber/v2"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/propagation"
	semconv "go.opentelemetry.io/otel/semconv/v1.25.0"
	"go.opentelemetry.io/otel/trace"
)

var Tracer = otel.Tracer("fiber_tracing_middleware")

// Config defines the config for middleware.
type middlewareConfig struct {
	Tracer                trace.Tracer
	TracerStartAttributes []trace.SpanStartOption
}

// ConfigDefault is the default config
var configDefault = middlewareConfig{
	Tracer: Tracer,
	TracerStartAttributes: []trace.SpanStartOption{
		trace.WithNewRoot(),
	},
}

// Helper function to set default values
func configDefaults(config ...middlewareConfig) middlewareConfig {
	// Return default config if nothing provided
	if len(config) < 1 {
		return configDefault
	}

	// Override default config
	cfg := config[0]

	if len(cfg.TracerStartAttributes) == 0 {
		cfg.TracerStartAttributes = configDefault.TracerStartAttributes
	}

	return cfg
}

// A helper implementing `propagation.TextMapCarrier`
type headersCarrier struct {
	*fiber.Ctx
}

func (h headersCarrier) Get(key string) string {
	c := h.Ctx
	headers := c.GetReqHeaders()
	if val, ok := headers[key]; ok {
		if len(val) > 0 {
			return val[len(val)-1]
		}
	}
	return ""
}

func (h headersCarrier) Set(key string, value string) {
	c := h.Ctx
	c.Set(key, value)
	return
}

func (h headersCarrier) Keys() []string {
	c := h.Ctx
	headers := c.GetRespHeaders()
	keys := make([]string, 0)
	for k, _ := range headers {
		keys = append(keys, k)
	}
	return keys
}

func new(config ...middlewareConfig) fiber.Handler {
	// Set default config
	cfg := configDefaults(config...)

	// Return new handler
	return func(c *fiber.Ctx) error {
		spanStartAttributes := []attr.KeyValue{
			semconv.HTTPMethod(c.Method()),
			semconv.HTTPTarget(string(c.Request().RequestURI())),
			semconv.HTTPURL(c.OriginalURL()),
			semconv.HTTPUserAgent(string(c.Request().Header.UserAgent())),
			semconv.HTTPRequestContentLength(c.Request().Header.ContentLength()),
			semconv.HTTPScheme(c.Protocol()),
			semconv.NetTransportTCP,
		}
		spanStartAttributes = append(spanStartAttributes, attr.ProcessStart()...)

		opts := []trace.SpanStartOption{
			trace.WithAttributes(spanStartAttributes...),
			trace.WithSpanKind(trace.SpanKindServer),
		}
		opts = append(opts, cfg.TracerStartAttributes...)

		// Init context using values from request headers if possible:
		headersCarrier := headersCarrier{c}
		ctx := c.UserContext()
		tctx := propagation.TraceContext{}
		ctx = tctx.Extract(ctx, headersCarrier)

		otelCtx, span := Tracer.Start(
			ctx,
			c.Method()+" "+c.OriginalURL(),
			opts...,
		)

		c.SetUserContext(otelCtx)

		defer span.End()

		err := c.Next()

		statusCode := c.Response().StatusCode()
		attrs := semconv.HTTPResponseStatusCode(statusCode)
		span.SetAttributes(attrs)

		// Return with trace context added to headers
		ctx = c.UserContext()
		tctx.Inject(ctx, headersCarrier)

		return err
	}
}