package promptkit_test

import (
	"context"
	"errors"
	"fmt"
	"reflect"
	"sync"
	"testing"
	"testing/fstest"
	"time"

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

func TestListIdentitiesBasicContract(t *testing.T) {
	engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{Main: &configkit.Input{
		Source: configkit.Mapping("listing", map[string]any{
			"profiles": map[string]any{"configured-profile": map[string]any{"model": "example-model"}},
			"prompts":  map[string]any{"example-prompt": promptDefinitionWithVersion("10", listingContractPrompt())},
		}),
	}}})
	if err != nil {
		t.Fatal(err)
	}

	profiles, err := engine.ListProfiles(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if !containsListingProfile(profiles, "configured-profile") {
		t.Fatalf("profile listing %#v does not include configured profile", profiles)
	}

	prompts, err := engine.ListPrompts(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	want := []promptkit.PromptIdentity{
		{ID: "example-prompt", Version: stringPointerForTest("10")},
	}
	if len(prompts) != len(want) {
		t.Fatalf("prompt listing = %#v, want %#v", prompts, want)
	}
	for i := range want {
		if prompts[i].ID != want[i].ID || prompts[i].Version == nil || *prompts[i].Version != *want[i].Version {
			t.Fatalf("prompt listing = %#v, want %#v", prompts, want)
		}
	}

	profile, err := engine.InspectProfile(context.Background(), "configured-profile")
	if err != nil || profile == nil || profile.ProfileID != "configured-profile" {
		t.Fatalf("listed profile could not be inspected: (%+v, %v)", profile, err)
	}
	prompt, err := engine.InspectPrompt(context.Background(), prompts[0].ID)
	if err != nil || prompt == nil || prompt.PromptID != prompts[0].ID || prompt.PromptVersion == nil ||
		*prompt.PromptVersion != *prompts[0].Version {
		t.Fatalf("listed prompt could not be inspected: (%+v, %v)", prompt, err)
	}
}

func TestListIdentitiesInvalidEngine(t *testing.T) {
	var nilEngine *promptkit.Engine
	var zeroEngine promptkit.Engine
	for _, engine := range []*promptkit.Engine{nilEngine, &zeroEngine} {
		if _, err := engine.ListProfiles(context.Background()); !errors.Is(err, promptkit.ErrInvalidConfig) {
			t.Errorf("ListProfiles on invalid engine: %v", err)
		}
		if _, err := engine.ListPrompts(context.Background()); !errors.Is(err, promptkit.ErrInvalidConfig) {
			t.Errorf("ListPrompts on invalid engine: %v", err)
		}
	}
}

func TestListIdentitiesPreserveContextErrorsWithoutPartialResults(t *testing.T) {
	engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{Main: &configkit.Input{
		Source: configkit.Mapping("listing", map[string]any{
			"profiles": map[string]any{"profile": map[string]any{}},
			"prompts":  map[string]any{"example": promptDefinitionWithVersion("v1", listingContractPrompt())},
		}),
	}}})
	if err != nil {
		t.Fatal(err)
	}

	canceled, cancel := context.WithCancel(context.Background())
	cancel()
	expired, expire := context.WithDeadline(context.Background(), time.Now().Add(-time.Second))
	defer expire()
	for _, test := range []struct {
		name string
		ctx  context.Context
		want error
	}{
		{name: "canceled", ctx: canceled, want: context.Canceled},
		{name: "expired", ctx: expired, want: context.DeadlineExceeded},
	} {
		t.Run(test.name, func(t *testing.T) {
			profiles, err := engine.ListProfiles(test.ctx)
			if !errors.Is(err, test.want) || len(profiles) != 0 {
				t.Errorf("ListProfiles = %#v, %v; want no results and %v", profiles, err, test.want)
			}
			prompts, err := engine.ListPrompts(test.ctx)
			if !errors.Is(err, test.want) || len(prompts) != 0 {
				t.Errorf("ListPrompts = %#v, %v; want no results and %v", prompts, err, test.want)
			}
		})
	}
}

func TestListIdentitiesUseEffectiveLayeredMembership(t *testing.T) {
	underlay := configkit.Input{Source: configkit.Mapping("consumer defaults", map[string]any{
		"collections": map[string]any{"shared": map[string]any{
			"profiles": map[string]any{
				"base":    map[string]any{"endpoint": "https://example.test/v1", "model": "base-model"},
				"child":   map[string]any{"extends": "base"},
				"overlay": map[string]any{"endpoint": "https://example.test/v1", "model": "old-model"},
			},
			"prompts": map[string]any{
				"prompt": promptDefinitionWithVersion("v1", listingContractPrompt()),
			},
		}},
		"prompts": map[string]any{
			"underlay-only": promptDefinitionWithVersion("v1", listingContractPrompt()),
		},
	})}
	main := configkit.Input{Source: configkit.Mapping("application overlay", map[string]any{
		"collections": map[string]any{"shared": map[string]any{"profiles": map[string]any{
			"overlay": map[string]any{"model": "effective-model"},
		}}},
		"prompts": map[string]any{
			"main-only": promptDefinitionWithVersion("v3", listingContractPrompt()),
		},
	})}
	engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{
		Underlays: []configkit.Underlay{{Inputs: []configkit.Input{underlay}}},
		Main:      &main,
	}})
	if err != nil {
		t.Fatal(err)
	}

	profiles, err := engine.ListProfiles(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	for _, id := range []string{
		"shared/base", "shared/child", "shared/overlay",
	} {
		if !containsListingProfile(profiles, id) {
			t.Errorf("profile listing %#v does not include %q", profiles, id)
		}
	}
	for _, id := range []string{"shared/base", "shared/child", "shared/overlay"} {
		if got := countListingProfile(profiles, id); got != 1 {
			t.Errorf("profile %q appears %d times, want 1", id, got)
		}
	}

	overlay, err := engine.InspectProfile(context.Background(), "shared/overlay")
	if err != nil || overlay == nil || overlay.EffectiveModelParams.Model != "effective-model" {
		t.Fatalf("overlay profile did not resolve to its effective definition: (%+v, %v)", overlay, err)
	}
	child, err := engine.InspectProfile(context.Background(), "shared/child")
	if err != nil || child == nil || child.EffectiveModelParams.Model != "base-model" {
		t.Fatalf("inherited child profile did not resolve its effective definition: (%+v, %v)", child, err)
	}

	prompts, err := engine.ListPrompts(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	wantPrompts := []promptkit.PromptIdentity{
		{ID: "main-only", Version: stringPointerForTest("v3")},
		{ID: "shared/prompt", Version: stringPointerForTest("v1")},
		{ID: "underlay-only", Version: stringPointerForTest("v1")},
	}
	if !reflect.DeepEqual(prompts, wantPrompts) {
		t.Fatalf("effective prompt identities = %#v, want %#v", prompts, wantPrompts)
	}
	if got := countListingPrompt(prompts, "shared/prompt", "v1"); got != 1 {
		t.Fatalf("overlaid prompt identity appears %d times, want 1", got)
	}
}

func TestListIdentitiesEmptyEffectiveConfiguration(t *testing.T) {
	engine, err := promptkit.NewEngine(promptkit.Config{})
	if err != nil {
		t.Fatal(err)
	}
	profiles, err := engine.ListProfiles(context.Background())
	if err != nil || len(profiles) != 0 {
		t.Fatalf("cleared profile listing = %#v, %v", profiles, err)
	}
	prompts, err := engine.ListPrompts(context.Background())
	if err != nil || len(prompts) != 0 {
		t.Fatalf("cleared prompt listing = %#v, %v", prompts, err)
	}
}

func TestListIdentitiesAreOwnedAndSafeForConcurrentUse(t *testing.T) {
	engine, err := newListingContractEngine(t, map[string]any{
		"profiles": map[string]any{"profile": map[string]any{"model": "partial-model"}},
		"prompts":  map[string]any{"prompt": promptDefinitionWithVersion("v1", listingContractPrompt())},
	})
	if err != nil {
		t.Fatal(err)
	}

	profiles, err := engine.ListProfiles(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	prompts, err := engine.ListPrompts(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	for i := range profiles {
		if profiles[i] == "profile" {
			profiles[i] = "caller mutation"
		}
	}
	for i := range prompts {
		if prompts[i].ID == "prompt" {
			prompts[i].ID = "caller mutation"
			if prompts[i].Version != nil {
				*prompts[i].Version = "caller mutation"
			}
		}
	}
	profiles, err = engine.ListProfiles(context.Background())
	if err != nil || !containsListingProfile(profiles, "profile") {
		t.Fatalf("profile result mutation escaped into engine: %#v, %v", profiles, err)
	}
	prompts, err = engine.ListPrompts(context.Background())
	if err != nil || countListingPrompt(prompts, "prompt", "v1") != 1 {
		t.Fatalf("prompt result mutation escaped into engine: %#v, %v", prompts, err)
	}

	const workers = 4
	const iterations = 25
	start := make(chan struct{})
	results := make(chan error, workers)
	var wait sync.WaitGroup
	wait.Add(workers)
	for worker := 0; worker < workers; worker++ {
		go func() {
			defer wait.Done()
			<-start
			for iteration := 0; iteration < iterations; iteration++ {
				profiles, err := engine.ListProfiles(context.Background())
				if err != nil {
					results <- fmt.Errorf("list profiles: %w", err)
					return
				}
				if len(profiles) == 0 {
					results <- errors.New("profile listing is empty")
					return
				}
				profiles[0] = "worker mutation"

				prompts, err := engine.ListPrompts(context.Background())
				if err != nil {
					results <- fmt.Errorf("list prompts: %w", err)
					return
				}
				if len(prompts) == 0 {
					results <- errors.New("prompt listing is empty")
					return
				}
				prompts[0].ID = "worker mutation"
				if prompts[0].Version != nil {
					*prompts[0].Version = "worker mutation"
				}

				if _, err := engine.InspectProfile(context.Background(), "profile"); err != nil {
					results <- fmt.Errorf("inspect profile: %w", err)
					return
				}
				if _, err := engine.InspectPrompt(context.Background(), "prompt"); err != nil {
					results <- fmt.Errorf("inspect prompt: %w", err)
					return
				}
			}
		}()
	}
	close(start)
	wait.Wait()
	close(results)
	for err := range results {
		t.Error(err)
	}
}

func TestListIdentitiesUseCapturedSourcesWithoutExecutionWork(t *testing.T) {
	files := fstest.MapFS{
		"config/listing.yaml": &fstest.MapFile{Data: []byte("profiles:\n  captured-partial:\n    model: partial-model\nprompts:\n  captured-prompt:\n    version: v1\n    messages:\n      - role: user\n        content: Hello.\n    output:\n      format: text\n      validation_mode: basic\n")},
	}
	inputs, err := promptkit.FilesystemInputs("application listing", files, "config")
	if err != nil {
		t.Fatal(err)
	}
	engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{
		Underlays: []configkit.Underlay{{Inputs: inputs}},
	}}, promptkit.WithLLMClient(rejectingListingLLMClient{}), promptkit.WithArtifactReader(rejectingListingArtifactReader{}))
	if err != nil {
		t.Fatal(err)
	}
	files["config/listing.yaml"].Data = []byte("profiles:\n  replacement-profile: {}\nprompts:\n  replacement-prompt:\n    version: v2\n    messages:\n      - role: user\n        content: Changed.\n    output:\n      format: text\n      validation_mode: basic\n")

	profiles, err := engine.ListProfiles(context.Background())
	if err != nil || !containsListingProfile(profiles, "captured-partial") || containsListingProfile(profiles, "replacement-profile") {
		t.Fatalf("profile listing did not use captured source: %#v, %v", profiles, err)
	}
	prompts, err := engine.ListPrompts(context.Background())
	if err != nil || countListingPrompt(prompts, "captured-prompt", "v1") != 1 || countListingPrompt(prompts, "replacement-prompt", "v2") != 0 {
		t.Fatalf("prompt listing did not use captured source: %#v, %v", prompts, err)
	}
}

func listingContractPrompt() map[string]any {
	return map[string]any{
		"messages": []any{map[string]any{"role": "user", "content": "Hello."}},
		"output":   map[string]any{"format": "text", "validation_mode": "basic"},
	}
}

func containsListingProfile(profiles []string, id string) bool {
	for _, profile := range profiles {
		if profile == id {
			return true
		}
	}
	return false
}

func countListingProfile(profiles []string, id string) int {
	count := 0
	for _, profile := range profiles {
		if profile == id {
			count++
		}
	}
	return count
}

func countListingPrompt(prompts []promptkit.PromptIdentity, id, version string) int {
	count := 0
	for _, prompt := range prompts {
		if prompt.ID == id && prompt.Version != nil && *prompt.Version == version {
			count++
		}
	}
	return count
}

func newListingContractEngine(t *testing.T, values map[string]any, options ...promptkit.Option) (*promptkit.Engine, error) {
	t.Helper()
	input := configkit.Input{Source: configkit.Mapping("listing contract", values)}
	return promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{Main: &input}}, options...)
}

type rejectingListingLLMClient struct{}

func (rejectingListingLLMClient) Generate(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
	return nil, errors.New("unexpected model generation")
}

type rejectingListingArtifactReader struct{}

func (rejectingListingArtifactReader) Read(context.Context, promptkit.ArtifactRef) (*promptkit.Artifact, error) {
	return nil, errors.New("unexpected artifact read")
}