moteva/server/internal/application/layer_separation_test.go

package application

import (
	"bytes"
	"context"
	"image"
	"image/color"
	"image/png"
	"testing"
	"time"

	"img_infinite_canvas/internal/domain/design"
	"img_infinite_canvas/internal/infrastructure/memory"
)

func TestRunEditElementsSeparatesBackgroundForegroundAndTextLayers(t *testing.T) {
	ctx := context.Background()
	repo := memory.NewProjectRepository()
	assets := newFakeAssetStorage()
	queue := &fakeJobQueue{}
	extractor := &countingTextExtractor{
		result: design.TextExtraction{
			Layers: []design.ExtractedTextLayer{
				{Text: "Hi", X: 0.5, Y: 0.05, Width: 0.2, Height: 0.1, FontSize: 18, FontFamily: "Inter", FontWeight: "700", Color: "#111111"},
			},
		},
	}
	remover := &fakeBackgroundRemover{result: design.BackgroundRemoval{
		Image:       testLayerSeparationForegroundPNG(t),
		ContentType: "image/png",
		Width:       4,
		Height:      4,
	}}
	service := NewDesignService(repo, nil, assets, nil, extractor)
	service.SetJobQueue(queue)
	service.SetBackgroundRemover(remover)

	source := testLayerSeparationSourceDataURL()
	project := design.Project{
		ID:        "project-layer-separation",
		Title:     "Layer separation",
		Status:    design.StatusReady,
		UpdatedAt: time.Now(),
		Nodes: []design.Node{{
			ID:      "image-1",
			Type:    design.NodeTypeImage,
			Title:   "Source",
			X:       100,
			Y:       120,
			Width:   400,
			Height:  400,
			Content: source,
			Status:  "success",
		}},
	}
	if err := repo.Save(ctx, project); err != nil {
		t.Fatal(err)
	}

	started, err := service.RunNodeActionAsync(ctx, project.ID, "image-1", design.NodeActionRequest{Action: "edit-elements"})
	if err != nil {
		t.Fatal(err)
	}
	if len(queue.jobs) != 1 {
		t.Fatalf("expected one queued job, got %#v", queue.jobs)
	}
	if len(started.Nodes) != 2 {
		t.Fatalf("expected source plus separated placeholder, got %#v", started.Nodes)
	}
	if started.Nodes[0].Status != "success" || started.Nodes[0].Content != source {
		t.Fatalf("expected source node to remain unchanged, got %#v", started.Nodes[0])
	}
	if started.Nodes[1].Type != design.NodeTypeFrame || started.Nodes[1].Status != "generating" || started.Nodes[1].LayerRole != "manual-frame" || started.Nodes[1].ParentID != "" || started.Nodes[1].FillColor != "transparent" {
		t.Fatalf("expected separated placeholder beside source, got %#v", started.Nodes[1])
	}
	if started.Nodes[1].Content != source {
		t.Fatalf("expected separated placeholder to preview source image, got %#v", started.Nodes[1])
	}
	task, err := service.GeneratorTask(ctx, project.ID, started.LastThreadID)
	if err != nil {
		t.Fatal(err)
	}
	if task.GeneratorName != layerSeparationGeneratorName || task.Status != "submitted" {
		t.Fatalf("expected submitted layer separation task, got %#v", task)
	}

	if err := service.ProcessJob(ctx, queue.jobs[0]); err != nil {
		t.Fatal(err)
	}
	saved, err := repo.Get(ctx, project.ID)
	if err != nil {
		t.Fatal(err)
	}
	if saved.Status != design.StatusReady {
		t.Fatalf("expected ready project, got %s", saved.Status)
	}
	if len(saved.Nodes) != 5 {
		t.Fatalf("expected original, background, frame, foreground, and text nodes, got %#v", saved.Nodes)
	}
	if saved.Nodes[0].Status != "success" || saved.Nodes[0].Content != source {
		t.Fatalf("expected original source to remain unchanged, got %#v", saved.Nodes[0])
	}
	background := saved.Nodes[1]
	if background.Type != design.NodeTypeImage || background.LayerRole != "clean-image" || background.Status != "success" {
		t.Fatalf("expected clean background node, got %#v", background)
	}
	if background.X != 532 || background.Y != 120 {
		t.Fatalf("expected clean background at separated result origin, got %#v", background)
	}
	frame := saved.Nodes[2]
	if frame.Type != design.NodeTypeFrame || frame.LayerRole != "manual-frame" || frame.Status != "success" {
		t.Fatalf("expected separated frame beside source, got %#v", frame)
	}
	if frame.X != 532 || frame.Y != 120 {
		t.Fatalf("expected frame beside source, got %#v", frame)
	}
	if frame.ParentID != "" || frame.FillColor != "transparent" {
		t.Fatalf("expected top-level transparent separated frame, got %#v", frame)
	}
	if background.ParentID != frame.ID {
		t.Fatalf("expected clean background to be the bottom layer inside the frame, got %#v", background)
	}
	foreground := saved.Nodes[3]
	if foreground.LayerRole != "foreground-image" || foreground.ParentID != frame.ID {
		t.Fatalf("expected separated foreground node, got %#v", foreground)
	}
	if foreground.X != 632 || foreground.Y != 220 || foreground.Width != 200 || foreground.Height != 200 {
		t.Fatalf("unexpected foreground geometry: %#v", foreground)
	}
	text := saved.Nodes[4]
	if text.LayerRole != "editable-text" || text.ParentID != frame.ID || text.Content != "Hi" {
		t.Fatalf("expected editable text node, got %#v", text)
	}

	task, err = service.GeneratorTask(ctx, project.ID, started.LastThreadID)
	if err != nil {
		t.Fatal(err)
	}
	if task.Status != "completed" {
		t.Fatalf("expected completed task, got %s", task.Status)
	}
	if len(task.Artifacts) != 3 {
		t.Fatalf("expected foreground, background, and text artifacts, got %#v", task.Artifacts)
	}
	if task.Artifacts[0].Metadata.Label != layerSeparationForegroundLabel || task.Artifacts[1].Metadata.Label != layerSeparationBackgroundLabel || task.Artifacts[2].Metadata.Label != layerSeparationTextLabel {
		t.Fatalf("unexpected artifact labels: %#v", task.Artifacts)
	}
	if got := task.Artifacts[0].Metadata.BBox; len(got) != 4 || got[0] != 1 || got[1] != 1 || got[2] != 3 || got[3] != 3 {
		t.Fatalf("unexpected foreground bbox: %#v", got)
	}
	if task.Artifacts[2].Type != "text" || !bytes.Contains([]byte(task.Artifacts[2].Content), []byte(`"text":"Hi"`)) {
		t.Fatalf("expected text render data artifact, got %#v", task.Artifacts[2])
	}
}

func TestSubmitLayerSeparationTaskCompletesStandaloneArtifacts(t *testing.T) {
	ctx := context.Background()
	assets := newFakeAssetStorage()
	extractor := &countingTextExtractor{
		result: design.TextExtraction{
			Layers: []design.ExtractedTextLayer{{Text: "Title", X: 0.1, Y: 0.1, Width: 0.3, Height: 0.1}},
		},
	}
	remover := &fakeBackgroundRemover{result: design.BackgroundRemoval{
		Image:       testLayerSeparationForegroundPNG(t),
		ContentType: "image/png",
		Width:       4,
		Height:      4,
	}}
	service := NewDesignService(nil, nil, assets, nil, extractor)
	service.SetBackgroundRemover(remover)

	task, err := service.SubmitGeneratorTask(ctx, design.GeneratorTaskSubmitRequest{
		ProjectID:     "project-standalone-layer-separation",
		GeneratorName: layerSeparationGeneratorName,
		InputArgs: design.GeneratorTaskInputArgs{
			ImageURL: testLayerSeparationSourceDataURL(),
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	completed := waitForStandaloneGeneratorTask(t, service, task.GeneratorTaskID)
	if completed.Status != "completed" {
		t.Fatalf("expected completed standalone task, got %#v", completed)
	}
	if completed.GeneratorName != layerSeparationGeneratorName || len(completed.Artifacts) != 3 {
		t.Fatalf("unexpected standalone task: %#v", completed)
	}
	if completed.Artifacts[0].Metadata.Label != layerSeparationForegroundLabel || completed.Artifacts[2].Metadata.Label != layerSeparationTextLabel {
		t.Fatalf("unexpected standalone artifacts: %#v", completed.Artifacts)
	}
}

func TestPlainLayerSeparationBackgroundUsesSolidColorWhenBackdropIsPlain(t *testing.T) {
	img := image.NewRGBA(image.Rect(0, 0, 24, 24))
	for y := 0; y < 24; y++ {
		for x := 0; x < 24; x++ {
			img.SetRGBA(x, y, color.RGBA{R: 188, G: 185, B: 178, A: 255})
		}
	}
	mask := image.NewAlpha(img.Bounds())
	rect := image.Rect(8, 4, 16, 22)
	for y := rect.Min.Y; y < rect.Max.Y; y++ {
		for x := rect.Min.X; x < rect.Max.X; x++ {
			img.SetRGBA(x, y, color.RGBA{R: 32, G: 26, B: 22, A: 255})
			mask.SetAlpha(x, y, color.Alpha{A: 255})
		}
	}

	bg, ok := plainLayerSeparationBackgroundImage(img, mask, rect)
	if !ok {
		t.Fatal("expected plain background branch")
	}
	center := bg.RGBAAt(12, 12)
	if center.R != 188 || center.G != 185 || center.B != 178 || center.A != 255 {
		t.Fatalf("expected solid background color, got %#v", center)
	}
}

func TestPlainLayerSeparationBackgroundSkipsComplexBackdrop(t *testing.T) {
	img := image.NewRGBA(image.Rect(0, 0, 24, 24))
	for y := 0; y < 24; y++ {
		for x := 0; x < 24; x++ {
			img.SetRGBA(x, y, color.RGBA{R: uint8(20 + x*8), G: uint8(40 + y*7), B: uint8(80 + (x+y)*3), A: 255})
		}
	}
	mask := image.NewAlpha(img.Bounds())
	rect := image.Rect(9, 6, 15, 20)
	for y := rect.Min.Y; y < rect.Max.Y; y++ {
		for x := rect.Min.X; x < rect.Max.X; x++ {
			mask.SetAlpha(x, y, color.Alpha{A: 255})
		}
	}

	if _, ok := plainLayerSeparationBackgroundImage(img, mask, rect); ok {
		t.Fatal("expected complex background to use repair branch")
	}
}

func TestLayerSeparationPosterKeepsComplexBackgroundAndExtractsSmallElements(t *testing.T) {
	ctx := context.Background()
	extractor := &countingTextExtractor{
		result: design.TextExtraction{
			Layers: []design.ExtractedTextLayer{{
				Text:       "SALE",
				X:          0.1,
				Y:          0.1,
				Width:      0.35,
				Height:     0.125,
				FontSize:   12,
				FontFamily: "Inter",
				FontWeight: "700",
				Color:      "#111111",
			}},
		},
	}
	remover := &fakeBackgroundRemover{result: design.BackgroundRemoval{
		Image:       testPosterForegroundPNG(t),
		ContentType: "image/png",
		Width:       80,
		Height:      80,
	}}
	service := NewDesignService(nil, nil, nil, nil, extractor)
	service.SetBackgroundRemover(remover)

	result, err := service.createLayerSeparationResult(ctx, "project-poster-layer-separation", design.Node{
		ID:      "poster",
		Type:    design.NodeTypeImage,
		Title:   "Poster",
		Content: testPosterSourceDataURL(),
		Status:  "success",
	}, design.NodeActionRequest{Action: "edit-elements"})
	if err != nil {
		t.Fatal(err)
	}
	if len(result.TextExtraction.Layers) != 1 {
		t.Fatalf("expected editable text extraction to remain, got %#v", result.TextExtraction.Layers)
	}
	if len(result.Foregrounds) != 1 {
		t.Fatalf("expected only the small decorative foreground to be extracted, got %#v", result.Foregrounds)
	}
	if got := result.Foregrounds[0].BBox; len(got) != 4 || got[0] != 54 || got[1] != 44 || got[2] != 66 || got[3] != 56 {
		t.Fatalf("unexpected decorative foreground bbox: %#v", got)
	}

	backgroundData, _, err := loadImageContent(ctx, result.BackgroundURL)
	if err != nil {
		t.Fatal(err)
	}
	background, _, err := image.Decode(bytes.NewReader(backgroundData))
	if err != nil {
		t.Fatal(err)
	}
	topLeft := color.NRGBAModel.Convert(background.At(2, 2)).(color.NRGBA)
	bottomRight := color.NRGBAModel.Convert(background.At(76, 76)).(color.NRGBA)
	if topLeft == bottomRight {
		t.Fatalf("expected complex poster background to be preserved, got solid-looking pixels %#v", topLeft)
	}
	textPixel := color.NRGBAModel.Convert(background.At(14, 12)).(color.NRGBA)
	if textPixel.R < 32 && textPixel.G < 32 && textPixel.B < 32 {
		t.Fatalf("expected text pixels to be removed from background, got %#v", textPixel)
	}
	ornamentPixel := color.NRGBAModel.Convert(background.At(59, 49)).(color.NRGBA)
	if ornamentPixel.R > 200 && ornamentPixel.G < 80 && ornamentPixel.B < 80 {
		t.Fatalf("expected extracted ornament to be removed from background, got %#v", ornamentPixel)
	}
	scenePixel := color.NRGBAModel.Convert(background.At(20, 48)).(color.NRGBA)
	if scenePixel.B < 170 {
		t.Fatalf("expected large poster scenery to stay on the background, got %#v", scenePixel)
	}
}

func waitForStandaloneGeneratorTask(t *testing.T, service *DesignService, taskID string) design.GeneratorTask {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		task, err := service.GeneratorTask(context.Background(), "", taskID)
		if err != nil {
			t.Fatal(err)
		}
		if task.Status == "completed" || task.Status == "failed" {
			return task
		}
		time.Sleep(10 * time.Millisecond)
	}
	task, _ := service.GeneratorTask(context.Background(), "", taskID)
	t.Fatalf("timed out waiting for standalone task: %#v", task)
	return design.GeneratorTask{}
}

func testLayerSeparationSourceDataURL() string {
	img := image.NewNRGBA(image.Rect(0, 0, 4, 4))
	for y := 0; y < 4; y++ {
		for x := 0; x < 4; x++ {
			img.SetNRGBA(x, y, color.NRGBA{R: 240, G: 240, B: 240, A: 255})
		}
	}
	for y := 1; y < 3; y++ {
		for x := 1; x < 3; x++ {
			img.SetNRGBA(x, y, color.NRGBA{R: 200, G: 20, B: 30, A: 255})
		}
	}
	return testPNGDataURLFromImage(img)
}

func testLayerSeparationForegroundPNG(t *testing.T) []byte {
	t.Helper()
	img := image.NewNRGBA(image.Rect(0, 0, 4, 4))
	for y := 1; y < 3; y++ {
		for x := 1; x < 3; x++ {
			img.SetNRGBA(x, y, color.NRGBA{R: 200, G: 20, B: 30, A: 255})
		}
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		t.Fatal(err)
	}
	return buf.Bytes()
}

func testPosterSourceDataURL() string {
	return testPNGDataURLFromImage(testPosterSourceImage())
}

func testPosterSourceImage() *image.NRGBA {
	img := image.NewNRGBA(image.Rect(0, 0, 80, 80))
	for y := 0; y < 80; y++ {
		for x := 0; x < 80; x++ {
			img.SetNRGBA(x, y, color.NRGBA{R: uint8(190 + x/3), G: uint8(186 + y/4), B: uint8(174 + (x+y)/8), A: 255})
		}
	}
	for y := 38; y < 70; y++ {
		for x := 8; x < 50; x++ {
			img.SetNRGBA(x, y, color.NRGBA{R: 64, G: 106, B: 204, A: 255})
		}
	}
	for y := 8; y < 18; y++ {
		for x := 8; x < 36; x++ {
			img.SetNRGBA(x, y, color.NRGBA{R: 17, G: 17, B: 17, A: 255})
		}
	}
	for y := 44; y < 56; y++ {
		for x := 54; x < 66; x++ {
			img.SetNRGBA(x, y, color.NRGBA{R: 225, G: 32, B: 44, A: 255})
		}
	}
	return img
}

func testPosterForegroundPNG(t *testing.T) []byte {
	t.Helper()
	img := image.NewNRGBA(image.Rect(0, 0, 80, 80))
	for y := 8; y < 18; y++ {
		for x := 8; x < 36; x++ {
			img.SetNRGBA(x, y, color.NRGBA{R: 17, G: 17, B: 17, A: 255})
		}
	}
	for y := 38; y < 70; y++ {
		for x := 8; x < 50; x++ {
			img.SetNRGBA(x, y, color.NRGBA{R: 64, G: 106, B: 204, A: 255})
		}
	}
	for y := 44; y < 56; y++ {
		for x := 54; x < 66; x++ {
			img.SetNRGBA(x, y, color.NRGBA{R: 225, G: 32, B: 44, A: 255})
		}
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		t.Fatal(err)
	}
	return buf.Bytes()
}