package router

import (
	"testing"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/propagation"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/sdk/trace/tracetest"
	"go.opentelemetry.io/otel/trace/noop"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
)

func TestARequestIsTracedIntoTheEngine(t *testing.T) {
	rec := tracetest.NewSpanRecorder()
	prevTP, prevProp := otel.GetTracerProvider(), otel.GetTextMapPropagator()
	otel.SetTracerProvider(sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(rec)))
	otel.SetTextMapPropagator(propagation.TraceContext{})
	t.Cleanup(func() { otel.SetTracerProvider(prevTP); otel.SetTextMapPropagator(prevProp) })

	e := newFakeEngine(t, "m")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
	// The caller's own trace, continued rather than replaced.
	const caller = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
	res := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}},
		map[string]string{"traceparent": caller, "X-Request-Id": "req-1"})
	if res.Code != 200 {
		t.Fatalf("status %d: %s", res.Code, res.Body.String())
	}

	spans := map[string]sdktrace.ReadOnlySpan{}
	for _, s := range rec.Ended() {
		spans[s.Name()] = s
	}
	root, ok := spans["inference /v1/chat/completions"]
	if !ok {
		t.Fatalf("no inference span among %v", spans)
	}
	if got := root.SpanContext().TraceID().String(); got != "4bf92f3577b34da6a3ce929d0e0e4736" {
		t.Errorf("the caller's trace was not continued: %s", got)
	}
	attrs := map[string]string{}
	for _, a := range root.Attributes() {
		attrs[string(a.Key)] = a.Value.Emit()
	}
	if attrs["llmgw.request_id"] != "req-1" || attrs["llmgw.model"] != "m" ||
		attrs["http.response.status_code"] != "200" || attrs["llmgw.endpoint"] != "a" {
		t.Errorf("inference span attributes %v", attrs)
	}
	for _, name := range []string{"admission", "schedule", "upstream"} {
		s, ok := spans[name]
		if !ok {
			t.Errorf("no %s span", name)
			continue
		}
		if s.Parent().SpanID() != root.SpanContext().SpanID() {
			t.Errorf("%s is not under the request's span", name)
		}
	}
	up := spans["upstream"]
	tp, _ := e.traceparent.Load().(string)
	if want := "00-" + up.SpanContext().TraceID().String() + "-" + up.SpanContext().SpanID().String() + "-01"; tp != want {
		t.Errorf("the engine got traceparent %q, expected %q", tp, want)
	}
}

func TestWithoutTracingTheEngineGetsNoTraceparent(t *testing.T) {
	// What an unconfigured process has. Set explicitly, since the default
	// global delegates for good to the first real provider another test set.
	prevTP, prevProp := otel.GetTracerProvider(), otel.GetTextMapPropagator()
	otel.SetTracerProvider(noop.NewTracerProvider())
	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator())
	t.Cleanup(func() { otel.SetTracerProvider(prevTP); otel.SetTextMapPropagator(prevProp) })
	e := newFakeEngine(t, "m")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
	if res := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}}, nil); res.Code != 200 {
		t.Fatalf("status %d", res.Code)
	}
	if tp, _ := e.traceparent.Load().(string); tp != "" {
		t.Errorf("traceparent %q sent with tracing off", tp)
	}
}