package promptkit_test

import (
	"context"
	"fmt"
	"testing"

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

var inspectedRequirements *promptkit.RequirementInspection

func requirementInspectionExtras(entries int) map[string]any {
	providers := make(map[string]any, entries)
	for index := range entries {
		providers[fmt.Sprintf("provider-%03d", index)] = map[string]any{
			"enabled": true,
			"labels":  []any{fmt.Sprintf("tier-%d", index%4), fmt.Sprintf("region-%d", index%3)},
			"limits":  map[string]any{"burst": index + 1, "period": "minute"},
			"routes":  []any{map[string]any{"weight": index + 1, "fallback": true}},
		}
	}
	return map[string]any{"provider_options": providers}
}

func BenchmarkRequirementInspection(b *testing.B) {
	profiles := map[string]any{
		"small": map[string]any{
			"model": "inspection-model", "extra_params": requirementInspectionExtras(2),
		},
		"large": map[string]any{
			"model": "inspection-model", "extra_params": requirementInspectionExtras(64),
		},
	}
	definition := contractPrompt("small", "Hello {{.name}} {{input \"transcript\"}}", []any{
		map[string]any{"name": "transcript", "required": true},
	}, "", noValidationOutput)
	engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{Main: &configkit.Input{
		Source: configkit.Mapping("requirement inspection benchmark", map[string]any{
			"prompts":  oneVersionPrompt("inspection", "1", definition),
			"profiles": profiles,
		}),
	}}})
	if err != nil {
		b.Fatal(err)
	}
	ctx := context.Background()
	for _, profileID := range []string{"small", "large"} {
		request := promptkit.RequirementInspectionRequest{
			PromptID: "inspection", ProfileID: profileID, Vars: map[string]string{"name": "Ada"},
		}
		inspection, err := engine.InspectRequirements(ctx, request)
		if err != nil || inspection.PromptID != "inspection" || inspection.PromptVersion != "1" ||
			inspection.ProfileID != profileID || !inspection.Complete || len(inspection.Inputs) != 1 ||
			inspection.Inputs[0].Name != "transcript" || inspection.Inputs[0].Kind != promptkit.RequirementRequired ||
			len(inspection.Variables) != 1 || inspection.Variables[0].Name != "name" || !inspection.Variables[0].Supplied {
			b.Fatalf("requirement inspection fixture for %q is invalid: inspection=%+v err=%v", profileID, inspection, err)
		}
		selectedRequest := request
		b.Run(profileID+"_extra", func(b *testing.B) {
			b.ReportAllocs()
			for b.Loop() {
				inspection, err := engine.InspectRequirements(ctx, selectedRequest)
				if err != nil {
					b.Fatal(err)
				}
				inspectedRequirements = inspection
			}
		})
	}
}