z6: fix Nano Banana metadata handling

2 files changed, 73 insertions, 8 deletions

diff --git a/internal/image/nanobanana.go b/internal/image/nanobanana.go
index 0a49ed7..c17ad1a 100644
--- a/internal/image/nanobanana.go
+++ b/internal/image/nanobanana.go
@@ -26,8 +26,6 @@ const (
 	DefaultNanoBananaTextModel = "gemini-2.5-flash"
 
 	nanoBananaAspectRatio = "4:3"
-	nanoBananaImageWidth  = 800
-	nanoBananaImageHeight = 600
 	nanoBananaDataPrefix  = "data:image/png;base64,"
 	nanoBananaSource      = "nanobanana"
 )
@@ -124,14 +122,25 @@ func (c *NanoBananaClient) Search(ctx context.Context, opts *SearchOptions) ([]S
 	if err != nil {
 		return nil, err
 	}
+	width, height, err := decodedImageDimensions(imageBytes)
+	if err != nil {
+		return nil, err
+	}
+
+	description := fmt.Sprintf("Generated educational image for %s", opts.Query)
+	if translatedWord != "" {
+		description = fmt.Sprintf("%s (%s)", description, translatedWord)
+	} else if opts.Translation != "" {
+		description = fmt.Sprintf("%s (%s)", description, strings.TrimSpace(opts.Translation))
+	}
 
 	result := SearchResult{
 		ID:           c.generateImageID(opts.Query),
 		URL:          dataURL,
 		ThumbnailURL: dataURL,
-		Width:        nanoBananaImageWidth,
-		Height:       nanoBananaImageHeight,
-		Description:  fmt.Sprintf("Generated educational image for %s (%s)", opts.Query, translatedWord),
+		Width:        width,
+		Height:       height,
+		Description:  description,
 		Attribution:  "Generated by Google Gemini Nano Banana",
 		Source:       nanoBananaSource,
 	}
@@ -167,11 +176,18 @@ func (c *NanoBananaClient) Download(ctx context.Context, url string) (io.ReadClo
 
 // GetAttribution returns attribution text for the generated image.
 func (c *NanoBananaClient) GetAttribution(result *SearchResult) string {
+	width := 0
+	height := 0
+	if result != nil {
+		width = result.Width
+		height = result.Height
+	}
+
 	attribution := "Image generated by Google Gemini Nano Banana\n\n"
 	attribution += fmt.Sprintf("Model: %s\n", c.modelName())
 	attribution += fmt.Sprintf("Text model: %s\n", c.textModelName())
 	attribution += fmt.Sprintf("Aspect ratio: %s\n", nanoBananaAspectRatio)
-	attribution += fmt.Sprintf("Size: %dx%d\n", nanoBananaImageWidth, nanoBananaImageHeight)
+	attribution += fmt.Sprintf("Size: %dx%d\n", width, height)
 	if result != nil && result.Description != "" {
 		attribution += fmt.Sprintf("Result: %s\n", result.Description)
 	}
@@ -517,6 +533,15 @@ func normalizePNG(imageBytes []byte, mimeType string) ([]byte, error) {
 	return buffer.Bytes(), nil
 }
 
+func decodedImageDimensions(imageBytes []byte) (int, int, error) {
+	cfg, _, err := image.DecodeConfig(bytes.NewReader(imageBytes))
+	if err != nil {
+		return 0, 0, fmt.Errorf("decode generated image dimensions: %w", err)
+	}
+
+	return cfg.Width, cfg.Height, nil
+}
+
 func (c *NanoBananaClient) generateImageID(word string) string {
 	hash := md5.Sum([]byte(word))
 	return hex.EncodeToString(hash[:])[:8]
diff --git a/internal/image/nanobanana_test.go b/internal/image/nanobanana_test.go
index 39ce208..84bb57b 100644
--- a/internal/image/nanobanana_test.go
+++ b/internal/image/nanobanana_test.go
@@ -85,6 +85,7 @@ func TestNanoBananaClient_Search_CustomPromptSkipsTextGeneration(t *testing.T) {
 
 	results, err := client.Search(context.Background(), &SearchOptions{
 		Query:        "ябълка",
+		Translation:  "banana",
 		CustomPrompt: " custom flashcard prompt ",
 	})
 	if err != nil {
@@ -105,6 +106,9 @@ func TestNanoBananaClient_Search_CustomPromptSkipsTextGeneration(t *testing.T) {
 	if !strings.Contains(results[0].Description, "ябълка") {
 		t.Fatalf("result description = %q, want it to mention the query", results[0].Description)
 	}
+	if !strings.Contains(results[0].Description, "banana") {
+		t.Fatalf("result description = %q, want translation metadata from custom prompt", results[0].Description)
+	}
 	if !strings.HasPrefix(results[0].URL, "data:image/png;base64,") {
 		t.Fatalf("expected PNG data URI, got %q", results[0].URL)
 	}
@@ -192,8 +196,8 @@ func TestNanoBananaClient_Search_GeneratedPromptFlow(t *testing.T) {
 	if result.Source != nanoBananaSource {
 		t.Fatalf("Source = %q, want %q", result.Source, nanoBananaSource)
 	}
-	if result.Width != nanoBananaImageWidth || result.Height != nanoBananaImageHeight {
-		t.Fatalf("Size = %dx%d, want %dx%d", result.Width, result.Height, nanoBananaImageWidth, nanoBananaImageHeight)
+	if result.Width != 1 || result.Height != 1 {
+		t.Fatalf("Size = %dx%d, want %dx%d", result.Width, result.Height, 1, 1)
 	}
 	if !strings.Contains(result.Description, "apple") {
 		t.Fatalf("Description = %q, want translated word", result.Description)
@@ -265,6 +269,9 @@ func TestNanoBananaClient_Search_TranslationFailureFallsBackToQuery(t *testing.T
 	if !sceneSawOriginalQuery {
 		t.Fatal("expected scene generation to use the original query after translation failure")
 	}
+	if !strings.Contains(results[0].Description, "ябълка") {
+		t.Fatalf("Description = %q, want original query in fallback description", results[0].Description)
+	}
 	if client.GetLastPrompt() == "" {
 		t.Fatal("expected last prompt to be recorded")
 	}
@@ -319,6 +326,39 @@ func TestNanoBananaClient_Search_ImageGenerationError(t *testing.T) {
 	}
 }
 
+func TestNanoBananaClient_Search_CustomPromptPreservesTranslationMetadata(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
+	}
+	nanoBananaGenerateImage = func(_ context.Context, _ *NanoBananaClient, _ string) ([]byte, string, error) {
+		return mustPNGBytes(t), "image/png", nil
+	}
+
+	client := NewNanoBananaClient(&NanoBananaConfig{APIKey: "test-key"})
+	results, err := client.Search(context.Background(), &SearchOptions{
+		Query:        "ябълка",
+		Translation:  "banana",
+		CustomPrompt: "flashcard prompt",
+	})
+	if err != nil {
+		t.Fatalf("Search() unexpected error: %v", err)
+	}
+	if len(results) != 1 {
+		t.Fatalf("expected 1 result, got %d", len(results))
+	}
+	if !strings.Contains(results[0].Description, "banana") {
+		t.Fatalf("result description = %q, want translation metadata", results[0].Description)
+	}
+}
+
 func TestNanoBananaClient_Search_InvalidOptions(t *testing.T) {
 	t.Parallel()