package promptkit_test

import (
	"context"
	"errors"
	"testing"

	"go.maximumdirect.net/promptkit"
)

func TestPreparedNotificationSelectionAndOrdering(t *testing.T) {
	for _, split := range []bool{false, true} {
		t.Run(map[bool]string{false: "run", true: "prepared"}[split], func(t *testing.T) {
			calls := 0
			engine := observationEngine(t, observationClientFunc(func(_ context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
				if calls != 1 {
					t.Fatalf("generation before notification: %d", calls)
				}
				return &promptkit.GenerateResponse{Content: "ok"}, nil
			}))
			req := promptkit.RunRequest{PromptID: "prompt", OnPrepared: func(_ context.Context, details promptkit.PreparedRun) {
				calls++
				if details.PromptID != "prompt" || len(details.Messages) == 0 {
					t.Fatalf("details=%+v", details)
				}
			}}
			if split {
				handle, err := engine.PrepareExecution(context.Background(), req)
				if err != nil {
					t.Fatal(err)
				}
				if calls != 0 {
					t.Fatal("preparation notified")
				}
				req.OnPrepared = func(context.Context, promptkit.PreparedRun) { t.Fatal("selection changed") }
				_, err = engine.RunPrepared(context.Background(), handle)
				if err != nil {
					t.Fatal(err)
				}
			} else if _, err := engine.Run(context.Background(), req); err != nil {
				t.Fatal(err)
			}
			if calls != 1 {
				t.Fatalf("notifications=%d", calls)
			}
		})
	}
}

func TestCancellationInPreparedNotificationPreventsGeneration(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	engine := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
		t.Fatal("generation after cancellation")
		return nil, nil
	}))
	_, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", OnPrepared: func(context.Context, promptkit.PreparedRun) { cancel() }})
	if !errors.Is(err, context.Canceled) || !errors.Is(err, promptkit.ErrLLMGenerate) {
		t.Fatalf("error=%v", err)
	}
}