package llm

import (
	"encoding/json"
	"runtime"
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
)

var prepareGenerationBenchmarkSink *PreparedGeneration

func BenchmarkPrepareGeneration(b *testing.B) {
	client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
		BaseURL: "https://api.example.test/v1",
		Model:   "benchmark-model",
	})
	if err != nil {
		b.Fatalf("construct benchmark client: %v", err)
	}

	prompt := domain.RenderedPrompt{
		Messages: []domain.RenderedMessage{{Role: "user", Content: "Summarize this request."}},
	}
	tests := []struct {
		name    string
		request domain.GenerateRequest
	}{
		{
			name: "no_extras",
			request: domain.GenerateRequest{
				Prompt: prompt,
				Target: domain.ExecutionTarget{Model: "benchmark-model"},
			},
		},
		{
			name: "small_extras",
			request: domain.GenerateRequest{
				Prompt: prompt,
				Target: domain.ExecutionTarget{Model: "benchmark-model", ExtraParams: map[string]any{
					"priority": "low",
					"attempt":  3,
					"enabled":  true,
				}},
			},
		},
		{
			name: "1MiB_string_extra",
			request: domain.GenerateRequest{
				Prompt: prompt,
				Target: domain.ExecutionTarget{Model: "benchmark-model", ExtraParams: map[string]any{
					"payload": strings.Repeat("x", 1<<20),
				}},
			},
		},
		{
			name: "nested_json_extras",
			request: domain.GenerateRequest{
				Prompt: prompt,
				Target: domain.ExecutionTarget{Model: "benchmark-model", ExtraParams: map[string]any{
					"options": map[string]any{
						"enabled": true,
						"limits":  []any{2, 4, 8},
						"tags":    []string{"summary", "offline"},
					},
				}},
			},
		},
		{
			name: "structured_output",
			request: domain.GenerateRequest{
				Prompt: prompt,
				Target: domain.ExecutionTarget{Model: "benchmark-model"},
				StructuredOutput: &domain.StructuredOutputSpec{
					Type: domain.StructuredOutputJSONSchema,
					JSONSchema: &domain.StructuredOutputJSONSpec{
						Name:   "summary",
						Strict: true,
						Schema: map[string]any{
							"type": "object",
							"properties": map[string]any{
								"summary": map[string]any{"type": "string"},
								"count":   map[string]any{"type": "integer"},
							},
							"required":             []string{"summary", "count"},
							"additionalProperties": false,
						},
					},
				},
			},
		},
	}

	for _, tc := range tests {
		b.Run(tc.name, func(b *testing.B) {
			prepared, err := client.PrepareGeneration(tc.request)
			if err != nil {
				b.Fatalf("check prepared request: %v", err)
			}
			if !json.Valid(prepared.payload) {
				b.Fatal("prepared request is not valid JSON")
			}

			b.ReportAllocs()
			b.SetBytes(int64(len(prepared.payload)))
			b.ResetTimer()
			for range b.N {
				prepared, err := client.PrepareGeneration(tc.request)
				if err != nil {
					b.Fatalf("prepare generation: %v", err)
				}
				prepareGenerationBenchmarkSink = prepared
			}
			b.StopTimer()
			runtime.KeepAlive(prepareGenerationBenchmarkSink)
		})
	}
}