package middleware

import (
	"net/http"
	"strings"

	"github.com/google/uuid"

	"github.com/jordan/composed5/pkg/httpcontext"
	"github.com/jordan/composed5/pkg/logging"
)

// Trace ID headers in priority order.
const (
	TraceIDHeader      = "X-Trace-ID"
	CloudTraceHeader   = "X-Cloud-Trace-Context"
)

// Tracing returns middleware that extracts or generates trace IDs.
//
// Checks headers in order:
//  1. X-Trace-ID - direct trace ID
//  2. X-Cloud-Trace-Context - GCP format "TRACE_ID/SPAN_ID;o=OPTIONS"
//  3. Generates a new UUID if none found
//
// The trace ID is stored in context via httpcontext.SetTraceID and
// logging.WithTraceID, and set in the X-Trace-ID response header.
//
// Usage:
//
//	r.Use(middleware.RequestID())
//	r.Use(middleware.Tracing())
//	r.Use(middleware.RequestLogger(logger))
func Tracing() func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			traceID := extractTraceID(r)

			if traceID == "" {
				traceID = uuid.New().String()
			}

			// Store in context
			ctx := httpcontext.SetTraceID(r.Context(), traceID)
			ctx = logging.WithTraceID(ctx, traceID)

			// Set response header
			w.Header().Set(TraceIDHeader, traceID)

			next.ServeHTTP(w, r.WithContext(ctx))
		})
	}
}

// extractTraceID tries to extract a trace ID from known headers.
func extractTraceID(r *http.Request) string {
	// X-Trace-ID takes priority
	if traceID := r.Header.Get(TraceIDHeader); traceID != "" {
		return traceID
	}

	// X-Cloud-Trace-Context format: "TRACE_ID/SPAN_ID;o=OPTIONS"
	if cloudTrace := r.Header.Get(CloudTraceHeader); cloudTrace != "" {
		if idx := strings.IndexByte(cloudTrace, '/'); idx > 0 {
			return cloudTrace[:idx]
		}
		if idx := strings.IndexByte(cloudTrace, ';'); idx > 0 {
			return cloudTrace[:idx]
		}
		return cloudTrace
	}

	return ""
}