package promptkit_test

import (
	"context"
	"errors"
	"testing"
	"testing/synctest"
	"time"

	"go.maximumdirect.net/promptkit"
)

func TestCoordinateUsesIndependentOperationBudgets(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		engine := newBackendCapacityEngine(t, observationClientFunc(func(_ context.Context, req promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
			if req.Prompt.SessionID == "continuation" {
				time.Sleep(2 * time.Second)
			}
			return &promptkit.GenerateResponse{Content: "answer", Usage: promptkit.TokenUsage{PromptTokens: 2}}, nil
		}), 1, capacityInt(0), nil)
		parent := context.Background()
		initial, err := engine.Run(parent, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
		if err != nil {
			t.Fatal(err)
		}
		var group promptkit.ExecutionGroup
		var child context.Context
		var details []promptkit.PreparedRun
		var requestEvents, coordinatorEvents []promptkit.ExecutionEvent
		start := time.Now()
		result, err := engine.Coordinate(parent, promptkit.CoordinateRequest{Start: initial.Exchange, AdditionalTurns: 2, Observer: promptkit.ComposeObservers(group.Observe, func(e promptkit.ExecutionEvent) { coordinatorEvents = append(coordinatorEvents, e) }), Next: func(ctx context.Context, turn promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
			if turn.Result != nil {
				if ctx.Err() != nil || child == nil || !errors.Is(child.Err(), context.Canceled) {
					t.Fatal("child was not cleaned up before callback")
				}
				if _, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second}); err != nil {
					t.Fatalf("assessment: %v", err)
				}
			}
			if turn.RemainingTurns == 0 {
				return nil, nil
			}
			return &promptkit.RunRequest{PromptID: "prompt", SessionID: "continuation", OperationTimeout: 3 * time.Second, Observer: func(e promptkit.ExecutionEvent) { requestEvents = append(requestEvents, e) }, OnPrepared: func(ctx context.Context, prepared promptkit.PreparedRun) {
				child = ctx
				details = append(details, prepared)
			}}, nil
		}})
		if err != nil || result.Outcome != promptkit.CoordinationStopped || time.Since(start) != 4*time.Second || len(details) != 2 || len(details[0].Messages) != 3 || len(details[1].Messages) != 5 {
			t.Fatalf("result=%+v error=%v details=%v elapsed=%v", result, err, len(details), time.Since(start))
		}
		summary := group.Snapshot()
		if len(requestEvents) != 6 || len(coordinatorEvents) != 6 || len(summary.Members) != 2 || summary.Attempts != 2 || summary.Usage.PromptTokens.Value != 4 || !summary.TerminalCoverageComplete {
			t.Fatalf("summary=%+v request=%d coordinator=%d", summary, len(requestEvents), len(coordinatorEvents))
		}
	})
}

func TestCoordinateTimeoutLeavesParentAndSiblingUsable(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		engine := newBackendCapacityEngine(t, observationClientFunc(func(ctx context.Context, req promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
			if req.Prompt.SessionID == "slow" {
				<-ctx.Done()
				return nil, ctx.Err()
			}
			return &promptkit.GenerateResponse{Content: "answer"}, nil
		}), 1, capacityInt(0), nil)
		parent := context.Background()
		initial, err := engine.Run(parent, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
		if err != nil {
			t.Fatal(err)
		}
		var details promptkit.PreparedRun
		calls := 0
		result, err := engine.Coordinate(parent, promptkit.CoordinateRequest{Start: initial.Exchange, AdditionalTurns: 1, Next: func(context.Context, promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
			calls++
			return &promptkit.RunRequest{PromptID: "prompt", SessionID: "slow", OperationTimeout: time.Second, OnPrepared: func(_ context.Context, p promptkit.PreparedRun) { details = p }}, nil
		}})
		if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || calls != 1 || result.Outcome != "" || result.Undelivered != nil || details.ParentRunID != initial.RunID || parent.Err() != nil {
			t.Fatalf("result=%+v error=%v callbacks=%d", result, err, calls)
		}
		sibling, err := engine.Run(parent, promptkit.RunRequest{PromptID: "prompt", ContinueFrom: initial.Exchange})
		if err != nil || sibling.Exchange == nil {
			t.Fatalf("sibling=%v error=%v", sibling, err)
		}
	})
}