gen2d/backend/internal/service/prompt_agent_test.go

package service

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestRunPromptAgent_Fallback(t *testing.T) {
	// 无 API key 时走模板回退路径
	output, err := RunPromptAgent(context.Background(), PromptAgentInput{
		Tags:      []string{"像素", "中世纪", "战士"},
		AssetType: "sprite",
		UserNote:  "持盾",
	})
	if err != nil {
		t.Fatalf("RunPromptAgent failed: %v", err)
	}

	if output.Prompt == "" {
		t.Fatal("expected non-empty prompt")
	}
	for _, section := range []string{"【主题】", "【风格】", "【技术】"} {
		if !strings.Contains(output.Prompt, section) {
			t.Errorf("output missing section %s: %s", section, output.Prompt)
		}
	}
}

func TestRunPromptAgent_Sprite(t *testing.T) {
	output, err := RunPromptAgent(context.Background(), PromptAgentInput{
		Tags:      []string{"像素", "中世纪", "战士"},
		AssetType: "sprite",
	})
	if err != nil {
		t.Fatalf("RunPromptAgent failed: %v", err)
	}
	if !strings.Contains(output.Prompt, "spritesheet") {
		t.Errorf("sprite output should mention spritesheet: %s", output.Prompt)
	}
}

func TestRunPromptAgent_Background(t *testing.T) {
	output, err := RunPromptAgent(context.Background(), PromptAgentInput{
		Tags:      []string{"森林", "暗黑"},
		AssetType: "background",
	})
	if err != nil {
		t.Fatalf("RunPromptAgent failed: %v", err)
	}
	if !strings.Contains(output.Prompt, "1920x1080") {
		t.Errorf("background output should mention 1920x1080: %s", output.Prompt)
	}
}

func TestRunPromptAgent_UI(t *testing.T) {
	output, err := RunPromptAgent(context.Background(), PromptAgentInput{
		Tags:      []string{"简约", "科幻"},
		AssetType: "ui",
	})
	if err != nil {
		t.Fatalf("RunPromptAgent failed: %v", err)
	}
	if !strings.Contains(output.Prompt, "九宫格") {
		t.Errorf("ui output should mention 九宫格: %s", output.Prompt)
	}
}

func TestRunPromptAgent_Animation(t *testing.T) {
	output, err := RunPromptAgent(context.Background(), PromptAgentInput{
		Tags:      []string{"火焰", "魔法"},
		AssetType: "animation",
	})
	if err != nil {
		t.Fatalf("RunPromptAgent failed: %v", err)
	}
	if !strings.Contains(output.Prompt, "帧") {
		t.Errorf("animation output should mention 帧: %s", output.Prompt)
	}
}

func TestRunPromptAgent_EmptyUserNote(t *testing.T) {
	output, err := RunPromptAgent(context.Background(), PromptAgentInput{
		Tags:      []string{"水彩"},
		AssetType: "sprite",
	})
	if err != nil {
		t.Fatalf("RunPromptAgent failed: %v", err)
	}
	if output.Prompt == "" {
		t.Fatal("expected non-empty prompt")
	}
}

func TestRunPromptAgent_SingleTag(t *testing.T) {
	output, err := RunPromptAgent(context.Background(), PromptAgentInput{
		Tags:      []string{"赛博朋克"},
		AssetType: "background",
	})
	if err != nil {
		t.Fatalf("RunPromptAgent with single tag failed: %v", err)
	}
	if output.Prompt == "" {
		t.Fatal("expected non-empty prompt")
	}
}

func TestRunPromptAgent_WithChatModel(t *testing.T) {
	// 启动一个 mock OpenAI 兼容 API
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/chat/completions" {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"【主题】一个像素战士\n【风格】像素风格\n【技术】spritesheet"}}]}`))
	}))
	defer server.Close()

	// 保存原始配置，测试后恢复
	orig := llmCfg
	defer func() { llmCfg = orig }()
	llmCfg.BaseURL = server.URL
	llmCfg.APIKey = "test-key"
	llmCfg.Model = "test-model"
	llmCfg.Temperature = 0.7
	llmCfg.MaxTokens = 512

	output, err := RunPromptAgent(context.Background(), PromptAgentInput{
		Tags:      []string{"像素", "战士"},
		AssetType: "sprite",
	})
	if err != nil {
		t.Fatalf("RunPromptAgent failed: %v", err)
	}

	if !strings.Contains(output.Prompt, "像素战士") {
		t.Errorf("expected LLM response in output, got: %s", output.Prompt)
	}
	if output.RawText != output.Prompt {
		t.Error("expected RawText == Prompt for non-streaming response")
	}
}

func TestChatCompletion_NonStream(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.Write([]byte(`{"choices":[{"message":{"content":"test response"}}]}`))
	}))
	defer server.Close()

	orig := llmCfg
	defer func() { llmCfg = orig }()
	llmCfg.BaseURL = server.URL
	llmCfg.APIKey = "key"
	llmCfg.Model = "m"
	llmCfg.Temperature = 0.5

	result, err := chatCompletion(context.Background(), []chatMessage{
		{Role: "user", Content: "hello"},
	})
	if err != nil {
		t.Fatalf("chatCompletion failed: %v", err)
	}
	if result != "test response" {
		t.Errorf("expected 'test response', got %q", result)
	}
}

func TestChatCompletion_Stream(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n"))
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\" World\"}}]}\n\n"))
		w.Write([]byte("data: [DONE]\n\n"))
	}))
	defer server.Close()

	orig := llmCfg
	defer func() { llmCfg = orig }()
	llmCfg.BaseURL = server.URL
	llmCfg.APIKey = "key"
	llmCfg.Model = "m"
	llmCfg.Temperature = 0.5

	result, err := chatCompletion(context.Background(), []chatMessage{
		{Role: "user", Content: "hi"},
	})
	if err != nil {
		t.Fatalf("chatCompletion stream failed: %v", err)
	}
	if result != "Hello World" {
		t.Errorf("expected 'Hello World', got %q", result)
	}
}

func TestCallLLMRefine_FallbackOnError(t *testing.T) {
	orig := llmCfg
	defer func() { llmCfg = orig }()
	llmCfg.BaseURL = "http://invalid-url.invalid"
	llmCfg.APIKey = "key"
	llmCfg.Model = "m"
	llmCfg.Temperature = 0.5

	// API 调用失败时应回退到模板生成
	output, err := callLLMRefine(context.Background(), "用户选择标签: 测试\n素材类型: sprite")
	if err != nil {
		t.Fatalf("fallback should not error: %v", err)
	}
	if !strings.Contains(output.Prompt, "【主题】") {
		t.Error("fallback output should contain 三段式 structure")
	}
}

// 单元测试

func TestBuildMetaPrompt(t *testing.T) {
	in := PromptAgentInput{
		Tags:      []string{"像素", "地牢"},
		AssetType: "sprite",
		UserNote:  "需要发光效果",
	}
	result := buildMetaPrompt(in)

	checks := []string{
		"2D 游戏素材提示词工程师",
		"三段式结构",
		"用户选择标签: 像素、地牢",
		"素材类型: sprite",
		"补充说明: 需要发光效果",
	}
	for _, c := range checks {
		if !strings.Contains(result, c) {
			t.Errorf("buildMetaPrompt missing %q", c)
		}
	}
}

func TestBuildMetaPrompt_NoUserNote(t *testing.T) {
	in := PromptAgentInput{
		Tags:      []string{"像素"},
		AssetType: "sprite",
	}
	result := buildMetaPrompt(in)
	if strings.Contains(result, "补充说明") {
		t.Error("buildMetaPrompt should not contain 补充说明 when UserNote is empty")
	}
}

func TestParseTagsFromMeta(t *testing.T) {
	meta := `用户选择标签: 像素、中世纪、战士
素材类型: sprite`
	tags, assetType := parseTagsFromMeta(meta)
	if len(tags) != 3 {
		t.Fatalf("expected 3 tags, got %d: %v", len(tags), tags)
	}
	if tags[0] != "像素" || tags[1] != "中世纪" || tags[2] != "战士" {
		t.Errorf("unexpected tags: %v", tags)
	}
	if assetType != "sprite" {
		t.Errorf("expected assetType=sprite, got %s", assetType)
	}
}

func TestParseTagsFromMeta_SingleTag(t *testing.T) {
	meta := `用户选择标签: 赛博朋克
素材类型: background`
	tags, _ := parseTagsFromMeta(meta)
	if len(tags) != 1 || tags[0] != "赛博朋克" {
		t.Errorf("expected [赛博朋克], got %v", tags)
	}
}

func TestParseTagsFromMeta_Empty(t *testing.T) {
	tags, assetType := parseTagsFromMeta("no tags here")
	if len(tags) != 0 || assetType != "" {
		t.Errorf("expected empty, got tags=%v assetType=%s", tags, assetType)
	}
}

func TestGenerateStructuredPrompt(t *testing.T) {
	prompt := generateStructuredPrompt([]string{"像素", "战士"}, "sprite")
	if !strings.HasPrefix(prompt, "【主题】") {
		t.Error("prompt should start with 【主题】")
	}
	themeIdx := strings.Index(prompt, "【主题】")
	styleIdx := strings.Index(prompt, "【风格】")
	techIdx := strings.Index(prompt, "【技术】")
	if !(themeIdx < styleIdx && styleIdx < techIdx) {
		t.Error("sections should be ordered: 主题 → 风格 → 技术")
	}
}

func TestBuildSubject(t *testing.T) {
	tags := []string{"像素", "战士"}
	tests := []struct {
		assetType, want string
	}{
		{"sprite", "精灵图"},
		{"background", "场景背景"},
		{"ui", "UI元素"},
		{"animation", "动画帧序列"},
		{"unknown", "游戏素材"},
	}
	for _, tt := range tests {
		result := buildSubject(tags, tt.assetType)
		if !strings.Contains(result, tt.want) {
			t.Errorf("buildSubject(%q) = %s, want containing %q", tt.assetType, result, tt.want)
		}
	}
}

func TestBuildStyle(t *testing.T) {
	result := buildStyle([]string{"像素", "暗黑"})
	if !strings.Contains(result, "色彩鲜明") {
		t.Error("style should contain default descriptions")
	}
	for _, tag := range []string{"像素", "暗黑"} {
		if !strings.Contains(result, tag+"风格") {
			t.Errorf("style missing %q", tag+"风格")
		}
	}
}

func TestBuildTechNotes(t *testing.T) {
	tests := []struct {
		assetType, want string
	}{
		{"sprite", "spritesheet"},
		{"background", "1920x1080"},
		{"ui", "九宫格"},
		{"animation", "4方向x4帧"},
		{"unknown", "PNG"},
	}
	for _, tt := range tests {
		result := buildTechNotes(tt.assetType)
		if !strings.Contains(result, tt.want) {
			t.Errorf("buildTechNotes(%q) = %s, want containing %q", tt.assetType, result, tt.want)
		}
	}
}