package trace_test

import (
	"context"
	"testing"

	"github.com/cloudwego/eino/compose"
	"github.com/hhs/camtalk/internal/trace"
)

func TestEinoContextPropagation(t *testing.T) {
	ctx := context.Background()
	testTraceID := "01J5TEST123456789"
	ctx = trace.WithTraceID(ctx, testTraceID)

	var capturedTraceID string

	g := compose.NewGraph[string, string]()
	g.AddLambdaNode("test_node", compose.InvokableLambda(
		func(ctx context.Context, input string) (string, error) {
			capturedTraceID = trace.GetTraceID(ctx)
			return "ok", nil
		},
	))
	g.AddEdge(compose.START, "test_node")
	g.AddEdge("test_node", compose.END)

	runnable, err := g.Compile(ctx)
	if err != nil {
		t.Fatalf("compile failed: %v", err)
	}

	_, err = runnable.Invoke(ctx, "test_input")
	if err != nil {
		t.Fatalf("invoke failed: %v", err)
	}

	if capturedTraceID != testTraceID {
		t.Errorf("trace_id lost in Eino propagation: got %q, want %q",
			capturedTraceID, testTraceID)
	}
}