gen2d/backend/tools/test_prompt_to_gif.go

//go:build ignore

package main

import (
	"bytes"
	"context"
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"image/gif"
	"image/png"
	"os"

	"gen2d/internal/config"
	"gen2d/internal/service"
	"gen2d/pkg/splitsprite"
)

func main() {
	cfg := config.Load()
	service.InitLLMConfig(cfg.LLM)
	service.InitImageGenConfig(cfg.ImageGen)

	ctx := context.Background()

	// 1. PromptAgent 优化提示词
	fmt.Println("=== Step 1: Optimize prompt via PromptAgent ===")
	agentIn := service.PromptAgentInput{
		Tags:      []string{"像素", "战士", "持剑", "精灵表"},
		AssetType: "sprite",
		Prompt:    "生成一个像素风持剑战士的4方向行走精灵表",
		UserNote:  "需要4方向（上下左右），每方向4帧行走动画",
	}
	out, err := service.RunPromptAgent(ctx, agentIn)
	if err != nil {
		fatalf("PromptAgent failed: %v", err)
	}
	fmt.Printf("Optimized prompt:\n%s\n\n", out.Prompt)

	// 2. 调用文生图 API
	fmt.Println("=== Step 2: Generate sprite sheet via image API ===")
	params := service.AssetParams{
		Resolution: 1024,
		Format:     "spritesheet",
	}
	images, err := service.GenerateImages(ctx, out.Prompt, params)
	if err != nil {
		fatalf("GenerateImages failed: %v", err)
	}
	if len(images) == 0 {
		fatalf("no images generated")
	}
	fmt.Printf("Generated %d image(s), size=%dx%d\n", len(images), images[0].Width, images[0].Height)

	os.MkdirAll("test_output", 0755)

	// 保存原始精灵表
	sheetPath := "test_output/sprite_sheet.png"
	if err := os.WriteFile(sheetPath, images[0].Data, 0644); err != nil {
		fatalf("save sheet: %v", err)
	}
	fmt.Printf("Saved sprite sheet → %s (%d bytes)\n", sheetPath, len(images[0].Data))

	// 3. splitsprite 拆分精灵表
	fmt.Println("\n=== Step 3: Split sprite sheet ===")
	sheetImg, err := decodePNG(images[0].Data)
	if err != nil {
		fatalf("decode sheet: %v", err)
	}
	opts := splitsprite.DefaultOptions()
	opts.GridRows = 4
	opts.GridCols = 4
	opts.GridPadding = 2
	opts.CenterAlign = true
	frames, err := splitsprite.Process(sheetImg, opts)
	if err != nil {
		fatalf("split failed: %v", err)
	}
	fmt.Printf("Detected %d frames\n", len(frames))

	// 保存单帧
	for i, f := range frames {
		fn := fmt.Sprintf("test_output/frame_%03d.png", i)
		if err := savePNG(fn, f); err != nil {
			fatalf("save frame %d: %v", i, err)
		}
	}
	fmt.Printf("Saved %d frames → test_output/frame_*.png\n", len(frames))

	// 4. 生成 GIF 预览
	fmt.Println("\n=== Step 4: Generate GIF preview ===")
	gifPath := "test_output/preview.gif"
	if err := genGIF(gifPath, frames, 12); err != nil {
		fatalf("generate GIF: %v", err)
	}
	fmt.Printf("GIF preview → %s (%d frames)\n", gifPath, len(frames))

	fmt.Println("\n=== Done ===")
	fmt.Println("Output files:")
	fmt.Println("  test_output/sprite_sheet.png  — original sprite sheet")
	fmt.Println("  test_output/frame_*.png       — individual frames")
	fmt.Println("  test_output/preview.gif       — animated GIF preview")
}

func genGIF(path string, frames []image.Image, delay int) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()

	pal := buildPalette(frames)
	anim := &gif.GIF{}
	for _, frame := range frames {
		b := frame.Bounds()
		paletted := image.NewPaletted(b, pal)
		draw.Draw(paletted, b, frame, b.Min, draw.Src)
		anim.Image = append(anim.Image, paletted)
		anim.Delay = append(anim.Delay, delay)
	}
	anim.LoopCount = 0 // loop forever
	return gif.EncodeAll(f, anim)
}

func buildPalette(frames []image.Image) color.Palette {
	hist := make(map[color.RGBA]int)
	sampleStep := max(1, len(frames)/8)
	for i := 0; i < len(frames); i += sampleStep {
		b := frames[i].Bounds()
		step := max(1, (b.Dx()*b.Dy())/4096)
		n := 0
		for y := b.Min.Y; y < b.Max.Y; y++ {
			for x := b.Min.X; x < b.Max.X; x++ {
				if n%step != 0 {
					n++
					continue
				}
				n++
				r, g, bl, a := frames[i].At(x, y).RGBA()
				if a > 0 {
					c := color.RGBA{R: uint8(r >> 8), G: uint8(g >> 8), B: uint8(bl >> 8), A: uint8(a >> 8)}
					hist[c]++
				}
			}
		}
	}
	pal := make(color.Palette, 0, 256)
	for c := range hist {
		pal = append(pal, c)
		if len(pal) >= 240 {
			break
		}
	}
	pal = append(pal,
		color.RGBA{0, 0, 0, 0},
		color.RGBA{0, 0, 0, 255},
		color.RGBA{255, 255, 255, 255},
	)
	return pal
}

func decodePNG(data []byte) (image.Image, error) {
	img, _, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return nil, err
	}
	return img, nil
}

func savePNG(path string, img image.Image) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()
	return png.Encode(f, img)
}

func fatalf(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", args...)
	os.Exit(1)
}