package promptkit_test

import (
	"context"
	"errors"
	"fmt"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/ConfigKit"
	"gitea.maximumdirect.net/eric/promptkit"
)

type observationClientFunc func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error)

func (f observationClientFunc) Generate(ctx context.Context, req promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
	return f(ctx, req)
}

func TestDelegatedOperationsIsolateObservations(t *testing.T) {
	for _, mode := range []string{"run", "observed run", "prepared", "observed prepared", "prepare", "prepare execution"} {
		t.Run(mode, func(t *testing.T) {
			type contextKey struct{}
			ctx, cancel := context.WithTimeout(context.WithValue(context.Background(), contextKey{}, "caller-value"), time.Minute)
			defer cancel()
			deadline, _ := ctx.Deadline()
			child, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{Main: &configkit.Input{
				Source: configkit.Mapping("child execution", map[string]any{
					"profiles": map[string]any{"profile": map[string]any{"endpoint": "https://example.invalid/v1", "model": "child-model"}},
					"prompts": map[string]any{"child": map[string]any{
						"version":         "1",
						"default_profile": "profile",
						"messages":        []any{map[string]any{"role": "user", "content": "child message"}},
						"output":          map[string]any{"format": "text", "validation_mode": "none"},
					}},
				}),
			}}}, promptkit.WithLLMClient(observationClientFunc(func(callCtx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
				gotDeadline, ok := callCtx.Deadline()
				if !ok || gotDeadline != deadline || callCtx.Value(contextKey{}) != "caller-value" || callCtx.Done() != ctx.Done() {
					t.Error("delegation lost caller context behavior")
				}
				return &promptkit.GenerateResponse{Content: "child", Usage: promptkit.TokenUsage{PromptTokens: 7}}, nil
			})))
			if err != nil {
				t.Fatal(err)
			}
			var childEvents []promptkit.ExecutionEvent
			var childResult *promptkit.RunResult
			outer := observationEngine(t, observationClientFunc(func(callCtx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
				req := promptkit.RunRequest{PromptID: "child"}
				if mode == "observed run" || mode == "observed prepared" {
					req.Observer = func(event promptkit.ExecutionEvent) { childEvents = append(childEvents, event) }
				}
				var callErr error
				switch mode {
				case "run", "observed run":
					childResult, callErr = child.Run(callCtx, req)
				case "prepared", "observed prepared", "prepare execution":
					var handle *promptkit.PreparedExecution
					handle, callErr = child.PrepareExecution(callCtx, req)
					if callErr == nil {
						defer handle.Discard()
						if mode != "prepare execution" {
							childResult, callErr = child.RunPrepared(callCtx, handle)
						}
					}
				case "prepare":
					_, callErr = child.Prepare(callCtx, req)
				}
				return &promptkit.GenerateResponse{Content: "outer", Usage: promptkit.TokenUsage{PromptTokens: 1}}, callErr
			}))
			var outerEvents []promptkit.ExecutionEvent
			result, err := outer.Run(ctx, promptkit.RunRequest{PromptID: "prompt", Observer: func(event promptkit.ExecutionEvent) { outerEvents = append(outerEvents, event) }})
			if err != nil {
				t.Fatal(err)
			}
			if len(outerEvents) != 3 {
				t.Fatalf("outer events=%+v", outerEvents)
			}
			for _, event := range outerEvents[1:] {
				if event.RunID != result.RunID || event.PromptID != "prompt" || event.Model != "model" || event.Usage.PromptTokens.Value != 1 {
					t.Fatalf("child contaminated outer event: %+v", event)
				}
			}
			if outerEvents[2].Attempts != 1 {
				t.Fatalf("outer attempts=%d", outerEvents[2].Attempts)
			}
			if childResult != nil && childResult.RunID == result.RunID {
				t.Fatal("delegated run reused outer run ID")
			}
			if mode == "observed run" || mode == "observed prepared" {
				if len(childEvents) != 3 || childEvents[0].Kind != promptkit.ExecutionStarted || childEvents[2].Kind != promptkit.ExecutionFinished || childEvents[2].Attempts != 1 || childEvents[2].Usage.PromptTokens.Value != 7 || childEvents[2].RunID != childResult.RunID {
					t.Fatalf("child lifecycle=%+v", childEvents)
				}
			} else if len(childEvents) != 0 {
				t.Fatalf("unobserved child events=%+v", childEvents)
			}
		})
	}
}

func TestConcurrentDelegationDoesNotShareRecorder(t *testing.T) {
	child := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
		return &promptkit.GenerateResponse{Content: "child"}, nil
	}))
	ids := make(chan string, 2)
	outer := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
		errs := make(chan error, 2)
		for range 2 {
			go func() {
				r, err := child.Run(ctx, promptkit.RunRequest{PromptID: "prompt"})
				if r != nil {
					ids <- r.RunID
				}
				errs <- err
			}()
		}
		return &promptkit.GenerateResponse{Content: "outer"}, errors.Join(<-errs, <-errs)
	}))
	var events []promptkit.ExecutionEvent
	r, err := outer.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", Observer: func(e promptkit.ExecutionEvent) { events = append(events, e) }})
	if err != nil {
		t.Fatal(err)
	}
	if len(events) != 3 || events[2].Attempts != 1 {
		t.Fatalf("outer events=%+v", events)
	}
	first, second := <-ids, <-ids
	if first == second || first == r.RunID || second == r.RunID {
		t.Fatal("delegated identities are not independent")
	}
}

type classifiedObservationTestError struct{ cause error }

func (e *classifiedObservationTestError) Error() string { return "injected failure" }
func (e *classifiedObservationTestError) Unwrap() error { return e.cause }

func TestInjectedAttemptAndTerminalFailureCategoriesAgree(t *testing.T) {
	for _, tc := range []struct {
		name  string
		cause error
		want  promptkit.ExecutionFailureKind
	}{
		{"nil success", nil, promptkit.ExecutionFailureMalformedResponse},
		{"malformed", promptkit.ErrMalformedResponse, promptkit.ExecutionFailureMalformedResponse},
		{"transport", promptkit.ErrTransport, promptkit.ExecutionFailureTransport},
		{"arbitrary", errors.New("transport failed"), promptkit.ExecutionFailureGeneration},
		{"canceled", errors.Join(promptkit.ErrTransport, context.Canceled), promptkit.ExecutionFailureCanceled},
		{"deadline", errors.Join(promptkit.ErrMalformedResponse, context.DeadlineExceeded), promptkit.ExecutionFailureDeadline},
	} {
		t.Run(tc.name, func(t *testing.T) {
			var injected error
			var typed *classifiedObservationTestError
			if tc.cause != nil {
				typed = &classifiedObservationTestError{cause: tc.cause}
				injected = fmt.Errorf("wrapped: %w", typed)
			}
			calls := 0
			engine := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
				calls++
				return nil, injected
			}))
			var events []promptkit.ExecutionEvent
			retries := 3
			result, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", TransientRetries: &retries, Observer: func(e promptkit.ExecutionEvent) { events = append(events, e) }})
			if result != nil || !errors.Is(err, promptkit.ErrLLMGenerate) || calls != 1 || len(events) != 3 {
				t.Fatalf("result=%+v err=%v calls=%d events=%+v", result, err, calls, events)
			}
			if events[1].Failure != tc.want || events[2].Failure != tc.want {
				t.Fatalf("attempt=%s terminal=%s want=%s", events[1].Failure, events[2].Failure, tc.want)
			}
			if injected != nil {
				var got *classifiedObservationTestError
				if !errors.Is(err, injected) || !errors.Is(err, tc.cause) || !errors.As(err, &got) || got != typed {
					t.Fatalf("lost injected identity: %v", err)
				}
			} else if !errors.Is(err, promptkit.ErrMalformedResponse) {
				t.Fatalf("nil success error=%v", err)
			}
		})
	}
}