package telemetry_test

import (
	"context"
	omni "git.yuntuo.hk/k2safe/go-omni"
	"git.yuntuo.hk/k2safe/go-omni/httpx"
	"git.yuntuo.hk/k2safe/go-omni/logx"
	"git.yuntuo.hk/k2safe/go-omni/telemetry"
	"github.com/redis/go-redis/v9"
	"go.opentelemetry.io/otel/propagation"
	sdkmetric "go.opentelemetry.io/otel/sdk/metric"
	"go.opentelemetry.io/otel/sdk/metric/metricdata"
	"go.opentelemetry.io/otel/sdk/trace/tracetest"
	"go.opentelemetry.io/otel/trace"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestHTTPTracePropagationMetricsAndRedisNil(t *testing.T) {
	exporter := tracetest.NewInMemoryExporter()
	reader := sdkmetric.NewManualReader()
	runtime, err := telemetry.New(telemetry.Config{Service: "test", SampleRatio: 1, Exporter: exporter, Reader: reader})
	if err != nil {
		t.Fatal(err)
	}
	defer runtime.Close()
	ctx := context.Background()
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !strings.HasPrefix(r.Header.Get("Traceparent"), "00-11111111111111111111111111111111-") {
			t.Error("missing propagated trace")
		}
		io.WriteString(w, "ok")
	}))
	defer upstream.Close()
	app := omni.New(omni.WithLogger(false))
	app.Use(func(next httpx.Handler) httpx.Handler {
		return func(c *httpx.Context) error {
			err := next(c)
			if logx.TraceID(c.Context()) != "11111111111111111111111111111111" {
				t.Error("outer access log lost trace correlation")
			}
			return err
		}
	})
	app.Use(runtime.Middleware())
	app.GET("/users/:id", func(c *httpx.Context) error {
		if logx.TraceID(c.Context()) != "11111111111111111111111111111111" {
			t.Error("log correlation")
		}
		req, _ := http.NewRequestWithContext(c.Context(), "GET", upstream.URL+"?secret=hidden", nil)
		resp, err := (telemetry.Transport{Runtime: runtime}).RoundTrip(req)
		if err != nil {
			return err
		}
		io.ReadAll(resp.Body)
		resp.Body.Close()
		if req.Header.Get("Traceparent") != "" {
			t.Error("caller headers mutated")
		}
		return c.OK("ok")
	})
	req := httptest.NewRequest("GET", "/users/privateUser?secret=hidden", nil)
	req.Header.Set("Traceparent", "00-11111111111111111111111111111111-2222222222222222-01")
	rec := httptest.NewRecorder()
	app.ServeHTTP(rec, req)
	if rec.Code != 200 {
		t.Fatal(rec.Code)
	}
	hook := telemetry.RedisHook{Runtime: runtime}
	if err = hook.ProcessHook(func(context.Context, redis.Cmder) error { return redis.Nil })(ctx, redis.NewStringCmd(ctx, "GET", "secret")); err != redis.Nil {
		t.Fatal("cache miss was swallowed", err)
	}
	if err = runtime.Flush(ctx); err != nil {
		t.Fatal(err)
	}
	spans := exporter.GetSpans()
	if len(spans) != 3 {
		t.Fatalf("got %d spans", len(spans))
	}
	var client, server bool
	for _, s := range spans {
		if strings.Contains(s.Name, "privateUser") || strings.Contains(s.Name, "secret") {
			t.Fatal(s.Name)
		}
		if s.SpanKind == trace.SpanKindServer {
			server = true
			if s.Name != "GET /users/:id" || s.Parent.SpanID().String() != "2222222222222222" {
				t.Fatal(s.Name, s.Parent)
			}
		}
		if s.SpanKind == trace.SpanKindClient {
			client = true
		}
	}
	if !client || !server {
		t.Fatal("missing client/server span")
	}
	var metrics metricdata.ResourceMetrics
	if err = reader.Collect(ctx, &metrics); err != nil {
		t.Fatal(err)
	}
	var count int64
	for _, scope := range metrics.ScopeMetrics {
		for _, m := range scope.Metrics {
			if sum, ok := m.Data.(metricdata.Sum[int64]); ok {
				for _, p := range sum.DataPoints {
					count += p.Value
				}
			}
		}
	}
	if count != 1 {
		t.Fatal("request metric", count)
	}
}

func TestConfiguredHeaderPropagation(t *testing.T) {
	exporter := tracetest.NewInMemoryExporter()
	propagator, err := telemetry.NewHeaderPropagator(telemetry.HeaderPropagationConfig{
		TraceIDHeader: "X-Trace-Id", SpanIDHeader: "X-Span-Id", RequestIDHeader: "X-Request-Id",
	})
	if err != nil {
		t.Fatal(err)
	}
	runtime, err := telemetry.New(telemetry.Config{Service: "header-test", SampleRatio: 1, Exporter: exporter, Propagator: propagator})
	if err != nil {
		t.Fatal(err)
	}
	defer runtime.Close()
	const traceID = "11111111111111111111111111111111"
	const parentSpanID = "2222222222222222"
	const requestID = "gateway-request-1"
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("X-Trace-Id") != traceID || r.Header.Get("X-Request-Id") != requestID {
			t.Errorf("custom propagation = trace %q request %q", r.Header.Get("X-Trace-Id"), r.Header.Get("X-Request-Id"))
		}
		if r.Header.Get("X-Span-Id") == parentSpanID || len(r.Header.Get("X-Span-Id")) != 16 || !strings.HasPrefix(r.Header.Get("Traceparent"), "00-"+traceID+"-") {
			t.Errorf("outbound parent propagation = %q", r.Header)
		}
		io.WriteString(w, "ok")
	}))
	defer upstream.Close()
	app := omni.New(omni.WithLogger(false), omni.WithRequestIDHeader("X-Request-Id"))
	defer app.Close()
	app.Use(runtime.Middleware())
	app.GET("/", func(c *httpx.Context) error {
		if got := logx.TraceID(c.Context()); got != traceID {
			t.Errorf("inbound trace correlation = %q", got)
		}
		if got := logx.RequestID(c.Context()); got != requestID {
			t.Errorf("inbound request correlation = %q", got)
		}
		req, _ := http.NewRequestWithContext(c.Context(), http.MethodGet, upstream.URL, nil)
		resp, err := (telemetry.Transport{Runtime: runtime}).RoundTrip(req)
		if err != nil {
			return err
		}
		defer resp.Body.Close()
		_, _ = io.ReadAll(resp.Body)
		return c.OK("ok")
	})
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.Header.Set("X-Trace-Id", traceID)
	req.Header.Set("X-Span-Id", parentSpanID)
	req.Header.Set("X-Request-Id", requestID)
	rec := httptest.NewRecorder()
	app.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK || rec.Header().Get("X-Request-Id") != requestID {
		t.Fatalf("response status/request ID = %d/%q", rec.Code, rec.Header().Get("X-Request-Id"))
	}
	if err := runtime.Flush(context.Background()); err != nil {
		t.Fatal(err)
	}
	spans := exporter.GetSpans()
	var server, client *tracetest.SpanStub
	for _, span := range spans {
		switch span.SpanKind {
		case trace.SpanKindServer:
			span := span
			server = &span
		case trace.SpanKindClient:
			span := span
			client = &span
		}
	}
	if server == nil || client == nil {
		t.Fatalf("missing spans: %d", len(spans))
	}
	if server.SpanContext.TraceID().String() != traceID || server.Parent.TraceID().String() != traceID || server.Parent.SpanID().String() != parentSpanID || !server.Parent.IsRemote() {
		t.Fatalf("server parent = %s/%s remote=%v", server.SpanContext.TraceID(), server.Parent.SpanID(), server.Parent.IsRemote())
	}
	if client.SpanContext.TraceID().String() != traceID || client.Parent.SpanID() != server.SpanContext.SpanID() {
		t.Fatalf("client linkage = trace %s parent %s", client.SpanContext.TraceID(), client.Parent.SpanID())
	}
}

func TestConfiguredHeaderPropagationPrefersW3CAndRejectsInvalidSetup(t *testing.T) {
	propagator, err := telemetry.NewHeaderPropagator(telemetry.HeaderPropagationConfig{TraceIDHeader: "x-trace-id", SpanIDHeader: "x-span-id"})
	if err != nil {
		t.Fatal(err)
	}
	carrier := propagation.MapCarrier{
		"traceparent": "00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-01",
		"X-Trace-Id":  "11111111111111111111111111111111",
		"X-Span-Id":   "2222222222222222",
	}
	ctx := propagator.Extract(context.Background(), carrier)
	parent := trace.SpanContextFromContext(ctx)
	if parent.TraceID().String() != "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" || parent.SpanID().String() != "bbbbbbbbbbbbbbbb" {
		t.Fatalf("W3C parent was not preferred: %s/%s", parent.TraceID(), parent.SpanID())
	}
	for _, config := range []telemetry.HeaderPropagationConfig{
		{TraceIDHeader: "", SpanIDHeader: "X-Span-Id"},
		{TraceIDHeader: "X-Trace-Id", SpanIDHeader: ""},
		{TraceIDHeader: "X-Trace-Id", SpanIDHeader: "X-Trace-Id"},
		{TraceIDHeader: "traceparent", SpanIDHeader: "X-Span-Id"},
		{TraceIDHeader: "X Trace", SpanIDHeader: "X-Span-Id"},
	} {
		if _, err := telemetry.NewHeaderPropagator(config); err == nil {
			t.Fatalf("invalid propagation config accepted: %+v", config)
		}
	}
	invalid := propagation.MapCarrier{"X-Trace-Id": "bad", "X-Span-Id": "bad"}
	if got := trace.SpanContextFromContext(propagator.Extract(context.Background(), invalid)); got.IsValid() {
		t.Fatal("invalid legacy IDs became a parent")
	}
}