package ngm

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"sync/atomic"
	"time"

	"github.com/gin-gonic/gin"
	"go.opentelemetry.io/otel/trace"
)

const traceIDHeader = "X-Trace-ID"

var (
	fallbackTraceIDCounter uint64
	fallbackSpanIDCounter  uint64
)

func traceIDHeaderMiddleware() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		traceID := ensureTraceID(ctx)
		if traceID != "" {
			ctx.Writer.Header().Set(traceIDHeader, traceID)
			ctx.Writer = &traceIDResponseWriter{
				ResponseWriter: ctx.Writer,
				traceID:        traceID,
			}
		}

		ctx.Next()

		traceID = traceIDFromContext(ctx.Request.Context())
		if traceID != "" {
			ctx.Writer.Header().Set(traceIDHeader, traceID)
		}
	}
}

type traceIDResponseWriter struct {
	gin.ResponseWriter
	traceID string
}

func (w *traceIDResponseWriter) WriteHeaderNow() {
	w.setTraceIDHeader()
	w.ResponseWriter.WriteHeaderNow()
}

func (w *traceIDResponseWriter) WriteHeader(code int) {
	w.setTraceIDHeader()
	w.ResponseWriter.WriteHeader(code)
}

func (w *traceIDResponseWriter) Write(data []byte) (int, error) {
	w.setTraceIDHeader()
	return w.ResponseWriter.Write(data)
}

func (w *traceIDResponseWriter) WriteString(data string) (int, error) {
	w.setTraceIDHeader()
	return w.ResponseWriter.WriteString(data)
}

func (w *traceIDResponseWriter) Flush() {
	w.setTraceIDHeader()
	w.ResponseWriter.Flush()
}

func (w *traceIDResponseWriter) setTraceIDHeader() {
	if w.traceID != "" {
		w.Header().Set(traceIDHeader, w.traceID)
	}
}

func ensureTraceID(ctx *gin.Context) string {
	traceID := traceIDFromContext(ctx.Request.Context())
	if traceID != "" {
		return traceID
	}

	spanCtx := trace.NewSpanContext(trace.SpanContextConfig{
		TraceID:    generateTraceID(),
		SpanID:     generateSpanID(),
		TraceFlags: trace.FlagsSampled,
	})
	ctx.Request = ctx.Request.WithContext(trace.ContextWithSpanContext(ctx.Request.Context(), spanCtx))
	return spanCtx.TraceID().String()
}

func traceIDFromContext(ctx context.Context) string {
	spanCtx := trace.SpanContextFromContext(ctx)
	if !spanCtx.HasTraceID() {
		return ""
	}
	return spanCtx.TraceID().String()
}

func generateTraceID() trace.TraceID {
	var traceID trace.TraceID
	if _, err := rand.Read(traceID[:]); err == nil && traceID.IsValid() {
		return traceID
	}
	binary.BigEndian.PutUint64(traceID[:8], uint64(time.Now().UnixNano()))
	binary.BigEndian.PutUint64(traceID[8:], atomic.AddUint64(&fallbackTraceIDCounter, 1))
	return traceID
}

func generateSpanID() trace.SpanID {
	var spanID trace.SpanID
	if _, err := rand.Read(spanID[:]); err == nil && spanID.IsValid() {
		return spanID
	}
	binary.BigEndian.PutUint64(spanID[:], uint64(time.Now().UnixNano())^atomic.AddUint64(&fallbackSpanIDCounter, 1))
	return spanID
}