package promptkit_test

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

	"go.maximumdirect.net/promptkit"
)

func TestRunOperationTimeoutAttribution(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
			<-ctx.Done()
			return nil, ctx.Err()
		}))
		var terminal promptkit.ExecutionEvent
		_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second,
			Observer: func(event promptkit.ExecutionEvent) {
				if event.Kind == promptkit.ExecutionFinished {
					terminal = event
				}
			}})
		if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || !errors.Is(err, promptkit.ErrLLMGenerate) || terminal.Failure != promptkit.ExecutionFailureDeadline {
			t.Fatalf("error=%v terminal=%+v", err, terminal)
		}
	})
}

func TestRunProviderDeadlineIsNotOperationTimeout(t *testing.T) {
	engine := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
		return nil, context.DeadlineExceeded
	}))
	_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Minute})
	if !errors.Is(err, context.DeadlineExceeded) || errors.Is(err, promptkit.ErrOperationTimeout) {
		t.Fatalf("provider deadline attribution: %v", err)
	}
}

func TestPreparedExecutionRetainsTimeoutWithoutIdleTimer(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
			<-ctx.Done()
			return nil, ctx.Err()
		}))
		req := promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second}
		handle, err := engine.PrepareExecution(context.Background(), req)
		if err != nil {
			t.Fatal(err)
		}
		req.OperationTimeout = 0
		time.Sleep(2 * time.Second)
		start := time.Now()
		_, err = engine.RunPrepared(context.Background(), handle)
		if time.Since(start) != time.Second || !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || !errors.Is(err, promptkit.ErrLLMGenerate) {
			t.Fatalf("elapsed=%v error=%v", time.Since(start), err)
		}
	})
}

func TestPreparedExecutionRespectsCallerDeadline(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
			<-ctx.Done()
			return nil, ctx.Err()
		}))
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		handle, err := engine.PrepareExecution(ctx, promptkit.RunRequest{PromptID: "prompt", OperationTimeout: 5 * time.Second})
		if err != nil {
			t.Fatal(err)
		}
		time.Sleep(time.Second)
		_, err = engine.RunPrepared(ctx, handle)
		if !errors.Is(err, context.DeadlineExceeded) || errors.Is(err, promptkit.ErrOperationTimeout) {
			t.Fatalf("caller deadline attribution: %v", err)
		}
	})
}