package openai

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/go-chi/chi/v5/middleware"
)

func traceIDViaMiddleware(req *http.Request) string {
	if req == nil {
		return requestTraceID(nil)
	}
	var got string
	h := middleware.RequestID(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
		got = requestTraceID(r)
	}))
	h.ServeHTTP(httptest.NewRecorder(), req)
	return got
}

func TestRequestTraceIDPriority(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/v1/chat/completions?__trace_id=query-trace", nil)
	req.Header.Set("X-Ds2-Test-Trace", "header-trace")
	got := traceIDViaMiddleware(req)
	if got != "query-trace" {
		t.Fatalf("expected query trace id to win, got %q", got)
	}
}

func TestRequestTraceIDHeaderFallback(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/v1/chat/completions", nil)
	req.Header.Set("X-Ds2-Test-Trace", "header-trace")
	got := traceIDViaMiddleware(req)
	if got != "header-trace" {
		t.Fatalf("expected header trace id to win when query missing, got %q", got)
	}
}

func TestRequestTraceIDReqIDFallback(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/v1/chat/completions", nil)
	got := traceIDViaMiddleware(req)
	if got == "" {
		t.Fatal("expected middleware request id fallback to be non-empty")
	}
}