z6: fix Nano Banana prompt and normalization paths

2 files changed, 187 insertions, 9 deletions

diff --git a/internal/image/nanobanana.go b/internal/image/nanobanana.go
index 6830b00..c65d8d4 100644
--- a/internal/image/nanobanana.go
+++ b/internal/image/nanobanana.go
@@ -8,6 +8,7 @@ import (
 	"encoding/hex"
 	"fmt"
 	"image"
+	_ "image/jpeg"
 	"image/png"
 	"io"
 	"net/http"
@@ -52,6 +53,12 @@ type NanoBananaClient struct {
 var _ ImageSearcher = (*NanoBananaClient)(nil)
 
 var newNanoBananaClient = genai.NewClient
+var nanoBananaGenerateText = func(ctx context.Context, c *NanoBananaClient, model, systemPrompt, userPrompt string, temperature float32, maxOutputTokens int32) (string, error) {
+	return c.generateText(ctx, model, systemPrompt, userPrompt, temperature, maxOutputTokens)
+}
+var nanoBananaGenerateImage = func(ctx context.Context, c *NanoBananaClient, prompt string) ([]byte, string, error) {
+	return c.generateImage(ctx, prompt)
+}
 
 // NewNanoBananaClient creates a new Nano Banana client.
 func NewNanoBananaClient(config *NanoBananaConfig) *NanoBananaClient {
@@ -88,12 +95,7 @@ func (c *NanoBananaClient) Search(ctx context.Context, opts *SearchOptions) ([]S
 		}
 	}
 
-	translatedWord, err := c.resolveTranslation(ctx, opts)
-	if err != nil {
-		return nil, err
-	}
-
-	prompt, err := c.resolvePrompt(ctx, opts, translatedWord)
+	prompt, translatedWord, err := c.buildPrompt(ctx, opts)
 	if err != nil {
 		return nil, err
 	}
@@ -106,7 +108,7 @@ func (c *NanoBananaClient) Search(ctx context.Context, opts *SearchOptions) ([]S
 	fmt.Printf("Nano Banana Image Generation Prompt (%d chars): %s\n", len(prompt), prompt)
 	fmt.Printf("Nano Banana Image Generation: Using model '%s' with aspect ratio '%s'\n", c.modelName(), nanoBananaAspectRatio)
 
-	imageBytes, mimeType, err := c.generateImage(ctx, prompt)
+	imageBytes, mimeType, err := nanoBananaGenerateImage(ctx, c, prompt)
 	if err != nil {
 		return nil, err
 	}
@@ -246,6 +248,36 @@ func (c *NanoBananaClient) resolvePrompt(ctx context.Context, opts *SearchOption
 	return c.createEducationalPrompt(ctx, opts.Query, translatedWord), nil
 }
 
+func (c *NanoBananaClient) buildPrompt(ctx context.Context, opts *SearchOptions) (string, string, error) {
+	if opts == nil {
+		return "", "", &SearchError{
+			Provider: nanoBananaSource,
+			Code:     "INVALID_OPTIONS",
+			Message:  "search options are required",
+		}
+	}
+
+	if customPrompt := strings.TrimSpace(opts.CustomPrompt); customPrompt != "" {
+		if len(customPrompt) > 1000 {
+			customPrompt = customPrompt[:997] + "..."
+		}
+		fmt.Printf("Using custom prompt: %s\n", customPrompt)
+		return customPrompt, "", nil
+	}
+
+	translatedWord, err := c.resolveTranslation(ctx, opts)
+	if err != nil {
+		return "", "", err
+	}
+
+	prompt, err := c.resolvePrompt(ctx, opts, translatedWord)
+	if err != nil {
+		return "", "", err
+	}
+
+	return prompt, translatedWord, nil
+}
+
 func (c *NanoBananaClient) createEducationalPrompt(ctx context.Context, bulgarianWord, englishTranslation string) string {
 	scene, err := c.generateSceneDescription(ctx, bulgarianWord, englishTranslation)
 	if err != nil {
@@ -317,8 +349,9 @@ func (c *NanoBananaClient) createEducationalPrompt(ctx context.Context, bulgaria
 func (c *NanoBananaClient) translateBulgarianToEnglish(ctx context.Context, word string) (string, error) {
 	fmt.Printf("Nano Banana Translation: Using model '%s' to translate '%s'\n", c.textModelName(), word)
 
-	translation, err := c.generateText(
+	translation, err := nanoBananaGenerateText(
 		ctx,
+		c,
 		c.textModelName(),
 		"You are a Bulgarian language expert. Translate the Bulgarian word into English. Respond with only the English translation, nothing else.",
 		fmt.Sprintf("Translate the Bulgarian word '%s' to English. Respond with only the English translation, nothing else.", word),
@@ -336,8 +369,9 @@ func (c *NanoBananaClient) translateBulgarianToEnglish(ctx context.Context, word
 func (c *NanoBananaClient) generateSceneDescription(ctx context.Context, bulgarianWord, englishTranslation string) (string, error) {
 	fmt.Printf("Nano Banana Scene Generation: Creating scene for '%s' (%s)\n", bulgarianWord, englishTranslation)
 
-	scene, err := c.generateText(
+	scene, err := nanoBananaGenerateText(
 		ctx,
+		c,
 		c.textModelName(),
 		"You are helping create educational flashcards for language learning. Generate a brief, vivid scene description that incorporates the given English word in a memorable, contextual way. The scene should be visually interesting and help with memory retention. Keep it to 1-2 sentences, focusing on visual elements that can be illustrated. The subject (the English word) should be the clear focal point of the image, prominent and centered.",
 		fmt.Sprintf("Create a scene description for the English word '%s' that would make a memorable flashcard image. Make sure '%s' is the main focus and most prominent element in the scene.", englishTranslation, englishTranslation),
diff --git a/internal/image/nanobanana_test.go b/internal/image/nanobanana_test.go
index 344dec9..925fcab 100644
--- a/internal/image/nanobanana_test.go
+++ b/internal/image/nanobanana_test.go
@@ -1,12 +1,19 @@
 package image
 
 import (
+	"bytes"
 	"context"
 	"encoding/base64"
+	"image"
+	"image/color"
+	"image/jpeg"
+	"image/png"
 	"io"
 	"os"
 	"strings"
 	"testing"
+
+	"google.golang.org/genai"
 )
 
 func TestNewNanoBananaClient(t *testing.T) {
@@ -47,6 +54,79 @@ func TestNanoBananaClient_NoAPIKey(t *testing.T) {
 	}
 }
 
+func TestNanoBananaClient_Search_CustomPromptSkipsTextGeneration(t *testing.T) {
+	originalText := nanoBananaGenerateText
+	originalImage := nanoBananaGenerateImage
+	t.Cleanup(func() {
+		nanoBananaGenerateText = originalText
+		nanoBananaGenerateImage = originalImage
+	})
+
+	nanoBananaGenerateText = func(_ context.Context, _ *NanoBananaClient, _, _, _ string, _ float32, _ int32) (string, error) {
+		t.Fatal("unexpected text generation for custom prompt")
+		return "", nil
+	}
+
+	var gotPrompt string
+	nanoBananaGenerateImage = func(_ context.Context, _ *NanoBananaClient, prompt string) ([]byte, string, error) {
+		gotPrompt = prompt
+		return mustPNGBytes(t), "image/png", nil
+	}
+
+	client := NewNanoBananaClient(&NanoBananaConfig{APIKey: "test-key"})
+	callbackCalled := false
+	client.SetPromptCallback(func(prompt string) {
+		callbackCalled = true
+		if prompt != "custom flashcard prompt" {
+			t.Fatalf("callback prompt = %q, want %q", prompt, "custom flashcard prompt")
+		}
+	})
+
+	results, err := client.Search(context.Background(), &SearchOptions{
+		Query:        "ябълка",
+		CustomPrompt: " custom flashcard prompt ",
+	})
+	if err != nil {
+		t.Fatalf("Search() unexpected error: %v", err)
+	}
+	if gotPrompt != "custom flashcard prompt" {
+		t.Fatalf("image prompt = %q, want %q", gotPrompt, "custom flashcard prompt")
+	}
+	if !callbackCalled {
+		t.Fatal("expected prompt callback to be called")
+	}
+	if client.GetLastPrompt() != "custom flashcard prompt" {
+		t.Fatalf("GetLastPrompt() = %q, want %q", client.GetLastPrompt(), "custom flashcard prompt")
+	}
+	if len(results) != 1 {
+		t.Fatalf("expected 1 result, got %d", len(results))
+	}
+	if !strings.Contains(results[0].Description, "ябълка") {
+		t.Fatalf("result description = %q, want it to mention the query", results[0].Description)
+	}
+	if !strings.HasPrefix(results[0].URL, "data:image/png;base64,") {
+		t.Fatalf("expected PNG data URI, got %q", results[0].URL)
+	}
+}
+
+func TestNanoBananaClient_Search_InvalidOptions(t *testing.T) {
+	t.Parallel()
+
+	client := NewNanoBananaClient(&NanoBananaConfig{APIKey: "test-key"})
+	_, err := client.Search(context.Background(), nil)
+	if err == nil {
+		t.Fatal("expected error for nil options")
+	}
+
+	searchErr, ok := err.(*SearchError)
+	if !ok {
+		t.Fatalf("expected SearchError, got %T", err)
+	}
+	if searchErr.Code != "INVALID_OPTIONS" {
+		t.Fatalf("expected INVALID_OPTIONS, got %s", searchErr.Code)
+	}
+}
+
 func TestNanoBananaClient_DownloadDataURI(t *testing.T) {
 	client := &NanoBananaClient{}
 	payload := []byte("png-bytes")
@@ -69,6 +149,42 @@ func TestNanoBananaClient_DownloadDataURI(t *testing.T) {
 	}
 }
 
+func TestNormalizePNG_FromJPEG(t *testing.T) {
+	t.Parallel()
+
+	jpegBytes := mustJPEGBytes(t)
+	pngBytes, err := normalizePNG(jpegBytes, "image/jpeg")
+	if err != nil {
+		t.Fatalf("normalizePNG() unexpected error: %v", err)
+	}
+	if !bytes.HasPrefix(pngBytes, []byte("\x89PNG\r\n\x1a\n")) {
+		t.Fatalf("normalizePNG() did not return PNG data")
+	}
+
+	img, err := png.Decode(bytes.NewReader(pngBytes))
+	if err != nil {
+		t.Fatalf("png.Decode() unexpected error: %v", err)
+	}
+	bounds := img.Bounds()
+	if bounds.Dx() != 1 || bounds.Dy() != 1 {
+		t.Fatalf("decoded PNG bounds = %v, want 1x1", bounds)
+	}
+}
+
+func TestExtractGeneratedImageErrors(t *testing.T) {
+	t.Parallel()
+
+	_, _, err := extractGeneratedImage(nil)
+	if err == nil {
+		t.Fatal("expected error for nil response")
+	}
+
+	_, _, err = extractGeneratedImage(&genai.GenerateContentResponse{})
+	if err == nil {
+		t.Fatal("expected error for empty response")
+	}
+}
+
 func TestNanoBananaClient_GetAttribution(t *testing.T) {
 	client := &NanoBananaClient{
 		config: &NanoBananaConfig{
@@ -125,3 +241,31 @@ func TestNanoBananaClient_Integration(t *testing.T) {
 func encodeBase64(data []byte) string {
 	return base64.StdEncoding.EncodeToString(data)
 }
+
+func mustPNGBytes(t *testing.T) []byte {
+	t.Helper()
+
+	img := image.NewRGBA(image.Rect(0, 0, 1, 1))
+	img.Set(0, 0, color.RGBA{R: 255, A: 255})
+
+	var buffer bytes.Buffer
+	if err := png.Encode(&buffer, img); err != nil {
+		t.Fatalf("png.Encode() unexpected error: %v", err)
+	}
+
+	return buffer.Bytes()
+}
+
+func mustJPEGBytes(t *testing.T) []byte {
+	t.Helper()
+
+	img := image.NewRGBA(image.Rect(0, 0, 1, 1))
+	img.Set(0, 0, color.RGBA{R: 0, G: 128, B: 255, A: 255})
+
+	var buffer bytes.Buffer
+	if err := jpeg.Encode(&buffer, img, &jpeg.Options{Quality: 90}); err != nil {
+		t.Fatalf("jpeg.Encode() unexpected error: %v", err)
+	}
+
+	return buffer.Bytes()
+}