package prompt

import (
	"bytes"
	"context"
	"runtime"
	"strconv"
	"testing"
	"text/template"

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

const rendererBenchmarkArtifactSize = 256 << 10

var (
	renderedPromptBenchmarkSink  *domain.RenderedPrompt
	renderedPromptsBenchmarkSink []*domain.RenderedPrompt
)

type rendererBenchmarkIdentity struct {
	id string
}

type rendererBenchmarkLookup struct {
	identity  rendererBenchmarkIdentity
	templates CompiledTemplates
}

func (lookup rendererBenchmarkLookup) GetCompiledTemplates(id string) (CompiledTemplates, bool) {
	if id != lookup.identity.id {
		return CompiledTemplates{}, false
	}
	return lookup.templates, true
}

type compiledRenderingBenchmarkCase struct {
	name        string
	renderer    Renderer
	definition  *domain.PromptDefinition
	inputs      map[string]*domain.Artifact
	vars        map[string]string
	ctx         context.Context
	outputBytes int
	check       func(*testing.B, *domain.RenderedPrompt)
}

func BenchmarkCompiledRendering(b *testing.B) {
	body := bytes.Repeat([]byte("x"), rendererBenchmarkArtifactSize)
	bodyText := string(body)

	small := newCompiledRenderingBenchmarkCase(b,
		"small_message", "", []string{"Summary for {{.audience}}."},
		&domain.PromptDefinition{
			ID: "compiled-render-benchmark", Version: stringPointerForTest("v1"),
			Templates: []domain.PromptMessageTemplate{{Role: "user", Content: "Summary for {{.audience}}."}},
		}, nil, map[string]string{"audience": "project leads"},
		len("Summary for project leads."),
		func(b *testing.B, rendered *domain.RenderedPrompt) {
			if rendered.SessionID != "" || len(rendered.Messages) != 1 || rendered.Messages[0].Content != "Summary for project leads." {
				b.Fatalf("unexpected small compiled render: %+v", rendered)
			}
		},
	)

	artifactInputs := map[string]*domain.Artifact{
		"document": {Name: "document", Body: body},
	}
	largeContent := `{{input "document"}}`
	largeDefinition := &domain.PromptDefinition{
		ID: "compiled-render-benchmark", Version: stringPointerForTest("v1"),
		Inputs:    []domain.PromptInput{{Name: "document"}},
		Templates: []domain.PromptMessageTemplate{{Role: "user", Content: largeContent}},
	}
	large := newCompiledRenderingBenchmarkCase(b,
		"artifact_backed_256KiB", "", []string{largeContent},
		largeDefinition, artifactInputs, nil, len(bodyText),
		func(b *testing.B, rendered *domain.RenderedPrompt) {
			if len(rendered.Messages) != 1 {
				b.Fatalf("artifact render message count=%d, want 1", len(rendered.Messages))
			}
			if rendered.Messages[0].Content != bodyText {
				b.Fatalf("artifact render length=%d, want %d", len(rendered.Messages[0].Content), len(bodyText))
			}
		},
	)

	repeatedSession := `document-size={{len (input "document")}}`
	repeatedMessages := []string{
		`{{input "document"}}`,
		`{{input "document"}}|{{input "document"}}`,
	}
	repeatedDefinition := &domain.PromptDefinition{
		ID: "compiled-render-benchmark", Version: stringPointerForTest("v1"), SessionID: repeatedSession,
		Inputs: []domain.PromptInput{{Name: "document"}},
		Templates: []domain.PromptMessageTemplate{
			{Role: "system", Content: repeatedMessages[0]},
			{Role: "user", Content: repeatedMessages[1]},
		},
	}
	repeatedOutputBytes := len(bodyText)*3 + 1
	repeated := newCompiledRenderingBenchmarkCase(b,
		"repeated_across_session_and_messages", repeatedSession, repeatedMessages,
		repeatedDefinition, artifactInputs, nil, repeatedOutputBytes,
		func(b *testing.B, rendered *domain.RenderedPrompt) {
			if rendered.SessionID != "document-size="+strconv.Itoa(len(bodyText)) || len(rendered.Messages) != 2 ||
				rendered.Messages[0].Content != bodyText || rendered.Messages[1].Content != bodyText+"|"+bodyText {
				b.Fatalf("unexpected repeated compiled render: session=%q message lengths=%v", rendered.SessionID, renderedMessageLengths(rendered))
			}
		},
	)

	for _, benchmark := range []compiledRenderingBenchmarkCase{small, large, repeated} {
		b.Run(benchmark.name, func(b *testing.B) {
			rendered, err := benchmark.renderer.Render(benchmark.ctx, benchmark.definition, benchmark.inputs, benchmark.vars)
			if err != nil {
				b.Fatalf("check compiled renderer output: %v", err)
			}
			benchmark.check(b, rendered)

			b.ReportAllocs()
			b.SetBytes(int64(benchmark.outputBytes))
			b.ResetTimer()
			for range b.N {
				rendered, err := benchmark.renderer.Render(benchmark.ctx, benchmark.definition, benchmark.inputs, benchmark.vars)
				if err != nil {
					b.Fatalf("render compiled template: %v", err)
				}
				renderedPromptBenchmarkSink = rendered
			}
			b.StopTimer()
			runtime.KeepAlive(renderedPromptBenchmarkSink)
		})
	}
}

// BenchmarkRetainedRenderedOutput reports post-GC live heap for a fixed batch
// of rendered outputs above the same fixture's setup-only baseline.
func BenchmarkRetainedRenderedOutput(b *testing.B) {
	body := bytes.Repeat([]byte("x"), rendererBenchmarkArtifactSize)
	bodyText := string(body)
	content := `{{input "document"}}`
	definition := &domain.PromptDefinition{
		ID: "retained-render-benchmark", Version: stringPointerForTest("v1"),
		Inputs:    []domain.PromptInput{{Name: "document"}},
		Templates: []domain.PromptMessageTemplate{{Role: "user", Content: content}},
	}
	compiled := CompiledTemplates{Messages: []*template.Template{mustCompileRendererBenchmarkTemplate(b, content)}}
	renderer := NewCompiledRenderer(rendererBenchmarkLookup{
		identity:  rendererBenchmarkIdentity{id: definition.ID},
		templates: compiled,
	})
	inputs := map[string]*domain.Artifact{"document": {Name: "document", Body: body}}
	ctx := context.Background()
	const retainedBatchSize = 4
	retained := make([]*domain.RenderedPrompt, retainedBatchSize)

	check, err := renderer.Render(ctx, definition, inputs, nil)
	if err != nil {
		b.Fatalf("check retained-output fixture: %v", err)
	}
	if len(check.Messages) != 1 || check.Messages[0].Content != bodyText {
		b.Fatalf("retained-output fixture returned unexpected content")
	}
	check = nil
	// Repeated -count runs must start from the same setup-only heap baseline.
	renderedPromptBenchmarkSink = nil
	renderedPromptsBenchmarkSink = nil

	runtime.GC()
	var setupMemory runtime.MemStats
	runtime.ReadMemStats(&setupMemory)
	runtime.KeepAlive(renderer)
	runtime.KeepAlive(definition)
	runtime.KeepAlive(inputs)

	b.ReportAllocs()
	b.SetBytes(int64(len(bodyText)))
	b.ResetTimer()
	for iteration := 0; iteration < b.N; iteration++ {
		rendered, err := renderer.Render(ctx, definition, inputs, nil)
		if err != nil {
			b.Fatalf("render retained-output fixture: %v", err)
		}
		retained[iteration%len(retained)] = rendered
	}
	b.StopTimer()
	for i, rendered := range retained {
		if rendered != nil {
			continue
		}
		rendered, err := renderer.Render(ctx, definition, inputs, nil)
		if err != nil {
			b.Fatalf("fill retained-output fixture: %v", err)
		}
		retained[i] = rendered
	}
	renderedPromptsBenchmarkSink = retained
	runtime.GC()
	var retainedMemory runtime.MemStats
	runtime.ReadMemStats(&retainedMemory)
	b.ReportMetric(float64(int64(retainedMemory.HeapAlloc)-int64(setupMemory.HeapAlloc)), "retained-B")
	runtime.KeepAlive(retained)
	runtime.KeepAlive(renderer)
	runtime.KeepAlive(definition)
	runtime.KeepAlive(inputs)
}

func newCompiledRenderingBenchmarkCase(
	b *testing.B,
	name string,
	session string,
	messages []string,
	definition *domain.PromptDefinition,
	inputs map[string]*domain.Artifact,
	vars map[string]string,
	outputBytes int,
	check func(*testing.B, *domain.RenderedPrompt),
) compiledRenderingBenchmarkCase {
	compiled := CompiledTemplates{}
	if session != "" {
		compiled.Session = mustCompileRendererBenchmarkTemplate(b, session)
	}
	for _, content := range messages {
		compiled.Messages = append(compiled.Messages, mustCompileRendererBenchmarkTemplate(b, content))
	}
	return compiledRenderingBenchmarkCase{
		name: name,
		renderer: NewCompiledRenderer(rendererBenchmarkLookup{
			identity:  rendererBenchmarkIdentity{id: definition.ID},
			templates: compiled,
		}),
		definition: definition, inputs: inputs, vars: vars, ctx: context.Background(),
		outputBytes: outputBytes, check: check,
	}
}

func mustCompileRendererBenchmarkTemplate(b *testing.B, source string) *template.Template {
	b.Helper()
	compiled, err := CompileTemplate("benchmark", source)
	if err != nil {
		b.Fatalf("compile benchmark template: %v", err)
	}
	return compiled
}

func renderedMessageLengths(rendered *domain.RenderedPrompt) []int {
	lengths := make([]int, len(rendered.Messages))
	for i, message := range rendered.Messages {
		lengths[i] = len(message.Content)
	}
	return lengths
}