path: root/internal/image/nanobanana.go

package image

import (
	"bytes"
	"context"
	"crypto/md5"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"image"
	_ "image/jpeg"
	"image/png"
	"io"
	"net/http"
	"strings"
	"time"

	"google.golang.org/genai"

	"codeberg.org/snonux/totalrecall/internal/apicircuit"
	"codeberg.org/snonux/totalrecall/internal/config"
	"codeberg.org/snonux/totalrecall/internal/httpctx"
)

const (
	// DefaultNanoBananaModel is the Gemini image model used for Nano Banana generation.
	DefaultNanoBananaModel = config.DefaultNanoBananaModel

	// DefaultNanoBananaTextModel is the Gemini text model used for translation and scene generation.
	DefaultNanoBananaTextModel = config.DefaultNanoBananaTextModel

	nanoBananaAspectRatio = "4:3"
	nanoBananaDataPrefix  = "data:image/png;base64,"
	nanoBananaSource      = "nanobanana"
)

// NanoBananaConfig holds the settings needed to build a Gemini-backed image generator.
type NanoBananaConfig struct {
	APIKey    string
	Model     string
	TextModel string
}

// NanoBananaClient implements ImageSearcher for Google Nano Banana image generation.
type NanoBananaClient struct {
	client     *genai.Client
	initErr    error
	config     *NanoBananaConfig
	lastPrompt string

	// PromptCallback is called when the prompt is generated, before the image is created.
	PromptCallback func(prompt string)
}

// Compile-time check that NanoBananaClient implements the full ImageClient interface
// (ImageSearcher + AttributionProvider).
var _ ImageClient = (*NanoBananaClient)(nil)

var newNanoBananaClient = httpctx.NewGenAIClient
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, aspectRatio string) ([]byte, string, error) {
	return c.generateImage(ctx, prompt, aspectRatio)
}

// NewNanoBananaClient creates a new Nano Banana client.
func NewNanoBananaClient(config *NanoBananaConfig) *NanoBananaClient {
	normalized := normalizeNanoBananaConfig(config)
	client := &NanoBananaClient{config: normalized}

	if normalized.APIKey == "" {
		return client
	}

	genaiClient, err := newNanoBananaClient(context.Background(), &genai.ClientConfig{
		APIKey:  normalized.APIKey,
		Backend: genai.BackendGeminiAPI,
	})
	if err != nil {
		client.initErr = err
		return client
	}

	client.client = genaiClient
	return client
}

// Search generates an educational image for the Bulgarian word using Nano Banana.
func (c *NanoBananaClient) Search(ctx context.Context, opts *SearchOptions) ([]SearchResult, error) {
	ctx, cancel := httpctx.WithTimeoutUnlessSet(ctx, httpctx.OperationTimeoutDefault)
	defer cancel()

	if err := c.ensureReady(); err != nil {
		return nil, err
	}
	if opts == nil {
		return nil, &SearchError{
			Provider: nanoBananaSource,
			Code:     "INVALID_OPTIONS",
			Message:  "search options are required",
		}
	}

	prompt, translatedWord, err := c.buildPrompt(ctx, opts)
	if err != nil {
		return nil, err
	}

	c.lastPrompt = prompt
	if c.PromptCallback != nil {
		c.PromptCallback(prompt)
	}

	// Resolve aspect ratio: use caller override if provided, else the default.
	aspectRatio := nanoBananaAspectRatio
	if opts.AspectRatio != "" {
		aspectRatio = opts.AspectRatio
	}

	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(), aspectRatio)

	var imageBytes []byte
	var mimeType string
	// Use multimodal chaining when reference images are available — this keeps
	// character appearance consistent across pages far more reliably than
	// injecting a text-only character bible.
	if len(opts.ReferenceImages) > 0 {
		imageBytes, mimeType, err = c.generateImageWithRefs(ctx, prompt, aspectRatio, opts.ReferenceImages)
	} else {
		imageBytes, mimeType, err = nanoBananaGenerateImage(ctx, c, prompt, aspectRatio)
	}
	if err != nil {
		if searchErr, ok := err.(*SearchError); ok {
			return nil, searchErr
		}
		return nil, &SearchError{
			Provider: nanoBananaSource,
			Code:     "API_ERROR",
			Message:  fmt.Sprintf("failed to generate image: %v", err),
		}
	}

	dataURL, err := encodeDataURL(imageBytes, mimeType)
	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)
	}

	result := SearchResult{
		ID:           c.generateImageID(opts.Query),
		URL:          dataURL,
		ThumbnailURL: dataURL,
		Width:        width,
		Height:       height,
		Description:  description,
		Attribution:  "Generated by Google Gemini Nano Banana",
		Source:       nanoBananaSource,
	}

	return []SearchResult{result}, nil
}

// Download returns the image bytes for either a data URI or a remote URL.
func (c *NanoBananaClient) Download(ctx context.Context, url string) (io.ReadCloser, error) {
	if strings.HasPrefix(url, nanoBananaDataPrefix) {
		return decodeDataURL(url)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}

	resp, err := imageHTTPClient.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		if closeErr := resp.Body.Close(); closeErr != nil {
			return nil, fmt.Errorf("HTTP %d: %s (failed to close response body: %v)", resp.StatusCode, resp.Status, closeErr)
		}
		return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, resp.Status)
	}

	return resp.Body, nil
}

// 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", width, height)
	if result != nil && result.Description != "" {
		attribution += fmt.Sprintf("Result: %s\n", result.Description)
	}
	attribution += fmt.Sprintf("\nPrompt used:\n%s\n", c.lastPrompt)
	attribution += fmt.Sprintf("\nGenerated at: %s\n", time.Now().Format("2006-01-02 15:04:05"))
	return attribution
}

// Name returns the name of the search provider.
func (c *NanoBananaClient) Name() string {
	return nanoBananaSource
}

// GetLastPrompt returns the last prompt used for image generation.
func (c *NanoBananaClient) GetLastPrompt() string {
	return c.lastPrompt
}

// SetPromptCallback sets a callback that runs after prompt generation.
func (c *NanoBananaClient) SetPromptCallback(callback func(prompt string)) {
	c.PromptCallback = callback
}

func (c *NanoBananaClient) ensureReady() error {
	if c == nil || c.config == nil {
		return &SearchError{
			Provider: nanoBananaSource,
			Code:     "NO_CONFIG",
			Message:  "Nano Banana client not initialized",
		}
	}
	if c.config.APIKey == "" {
		return &SearchError{
			Provider: nanoBananaSource,
			Code:     "NO_API_KEY",
			Message:  "Google API key not configured",
		}
	}
	if c.initErr != nil {
		return &SearchError{
			Provider: nanoBananaSource,
			Code:     "CLIENT_INIT_FAILED",
			Message:  fmt.Sprintf("failed to initialize client: %v", c.initErr),
		}
	}
	if c.client == nil {
		return &SearchError{
			Provider: nanoBananaSource,
			Code:     "CLIENT_NOT_READY",
			Message:  "Nano Banana client not initialized",
		}
	}

	return nil
}

// resolveTranslation returns the English translation for the Bulgarian query.
// Translating internally would couple the image package to the Gemini text API for a
// concern that belongs in the translation package. Callers (processor, GUI) already
// resolve the English translation before calling Search, so we simply use what was
// provided and fall back to the original query word when nothing was given.
func (c *NanoBananaClient) resolveTranslation(_ context.Context, opts *SearchOptions, translation string) (string, error) {
	if translation != "" {
		fmt.Printf("Using provided translation: %s -> %s\n", opts.Query, translation)
		return translation, nil
	}

	// No translation provided — fall back to the original query word so image
	// generation still proceeds, albeit potentially with lower quality.
	return opts.Query, nil
}

func (c *NanoBananaClient) resolvePrompt(ctx context.Context, opts *SearchOptions, translatedWord string) (string, error) {
	if customPrompt := strings.TrimSpace(opts.CustomPrompt); customPrompt != "" {
		if len(customPrompt) > 4000 {
			customPrompt = customPrompt[:3997] + "..."
		}
		fmt.Printf("Using custom prompt: %s\n", customPrompt)
		return customPrompt, nil
	}

	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",
		}
	}

	translation := strings.TrimSpace(opts.Translation)

	if customPrompt := strings.TrimSpace(opts.CustomPrompt); customPrompt != "" {
		if len(customPrompt) > 4000 {
			customPrompt = customPrompt[:3997] + "..."
		}
		fmt.Printf("Using custom prompt: %s\n", customPrompt)
		return customPrompt, translation, nil
	}

	translatedWord, err := c.resolveTranslation(ctx, opts, translation)
	if err != nil {
		return "", "", err
	}

	prompt, err := c.resolvePrompt(ctx, opts, translatedWord)
	if err != nil {
		return "", "", err
	}

	return prompt, translatedWord, nil
}

// createEducationalPrompt generates a prompt optimized for language learning.
// Scene generation and style selection are handled here; the shared
// buildEducationalPrompt helper assembles the actual prompt text so that the
// same policy is used by both NanoBananaClient and OpenAIClient.
func (c *NanoBananaClient) createEducationalPrompt(ctx context.Context, bulgarianWord, englishTranslation string) string {
	subject := promptSubject(englishTranslation, bulgarianWord)

	scene, err := c.generateSceneDescription(ctx, bulgarianWord, englishTranslation)
	if err != nil {
		fmt.Printf("  Failed to generate scene: %v, using basic prompt\n", err)
		scene = ""
	}
	if scene != "" {
		scene = sanitizeSceneDescription(scene)
		if !usableSceneDescription(scene) {
			fmt.Printf("  Scene response was too short or generic, using basic prompt\n")
			scene = ""
		}
	}

	// Select a random style from the shared pool. Fall back to a generic style
	// if the pool has been exhausted by tests or other callers.
	selectedStyle := chooseArtisticStyle()
	if selectedStyle == defaultArtisticStyle {
		fmt.Printf("  No artistic styles available, using generic prompt\n")
	}
	fmt.Printf("  Using image style: %s\n", selectedStyle)

	return buildEducationalPrompt(selectedStyle, scene, subject)
}

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 := 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),
		0.7,
		100,
	)
	if err != nil {
		return "", fmt.Errorf("scene generation failed: %w", err)
	}
	scene = sanitizeSceneDescription(scene)
	if !usableSceneDescription(scene) {
		return "", fmt.Errorf("scene generation returned unusable content")
	}

	fmt.Printf("Generated scene: %s\n", scene)
	return scene, nil
}

func (c *NanoBananaClient) generateText(ctx context.Context, model, systemPrompt, userPrompt string, temperature float32, maxOutputTokens int32) (string, error) {
	temp := temperature
	resp, err := apicircuit.GeminiNanoBanana(func() (*genai.GenerateContentResponse, error) {
		return c.client.Models.GenerateContent(ctx, model, []*genai.Content{
			genai.NewContentFromText(userPrompt, genai.RoleUser),
		}, &genai.GenerateContentConfig{
			SystemInstruction: genai.NewContentFromText(systemPrompt, genai.RoleUser),
			Temperature:       &temp,
			MaxOutputTokens:   maxOutputTokens,
		})
	})
	if err != nil {
		return "", fmt.Errorf("gemini API error: %w", err)
	}

	text := strings.TrimSpace(resp.Text())
	if text == "" {
		return "", fmt.Errorf("no response received")
	}

	return text, nil
}

func (c *NanoBananaClient) generateImage(ctx context.Context, prompt, aspectRatio string) ([]byte, string, error) {
	if aspectRatio == "" {
		aspectRatio = nanoBananaAspectRatio
	}
	cfg := &genai.GenerateContentConfig{
		ResponseModalities: []string{string(genai.ModalityImage)},
		ImageConfig: &genai.ImageConfig{
			AspectRatio: aspectRatio,
		},
	}

	resp, err := apicircuit.GeminiNanoBanana(func() (*genai.GenerateContentResponse, error) {
		return c.client.Models.GenerateContent(ctx, c.modelName(), []*genai.Content{
			genai.NewContentFromText(prompt, genai.RoleUser),
		}, cfg)
	})
	if err != nil {
		return nil, "", &SearchError{
			Provider: nanoBananaSource,
			Code:     "API_ERROR",
			Message:  fmt.Sprintf("failed to generate image: %v", err),
		}
	}

	imageBytes, mimeType, err := extractGeneratedImage(resp)
	if err != nil {
		return nil, "", &SearchError{
			Provider: nanoBananaSource,
			Code:     "NO_RESULTS",
			Message:  err.Error(),
		}
	}

	return imageBytes, mimeType, nil
}

// generateImageWithRefs sends reference images alongside the text prompt so the
// model can match character appearance from the existing pages. This implements
// the iterative chaining technique: each new page is conditioned on the visual
// look established by previous pages rather than relying on text descriptions alone.
func (c *NanoBananaClient) generateImageWithRefs(ctx context.Context, prompt, aspectRatio string, refs [][]byte) ([]byte, string, error) {
	if aspectRatio == "" {
		aspectRatio = nanoBananaAspectRatio
	}
	cfg := &genai.GenerateContentConfig{
		ResponseModalities: []string{string(genai.ModalityImage)},
		ImageConfig:        &genai.ImageConfig{AspectRatio: aspectRatio},
	}

	// Build multimodal content: reference image bytes first, then the instruction
	// + prompt text. Leading with images means they are processed before the text.
	parts := make([]*genai.Part, 0, len(refs)+1)
	for _, ref := range refs {
		if len(ref) > 0 {
			parts = append(parts, &genai.Part{
				InlineData: &genai.Blob{MIMEType: "image/png", Data: ref},
			})
		}
	}
	refNote := fmt.Sprintf(
		"The %d reference image(s) above show the EXACT character appearance that must be preserved. "+
			"Every character — same face, same age, same hair, same clothing, same animal breed and markings — "+
			"must look IDENTICAL in the new image. Now generate:\n\n",
		len(refs),
	)
	parts = append(parts, &genai.Part{Text: refNote + prompt})

	resp, err := apicircuit.GeminiNanoBanana(func() (*genai.GenerateContentResponse, error) {
		return c.client.Models.GenerateContent(ctx, c.modelName(),
			[]*genai.Content{{Role: string(genai.RoleUser), Parts: parts}},
			cfg,
		)
	})
	if err != nil {
		return nil, "", &SearchError{
			Provider: nanoBananaSource,
			Code:     "API_ERROR",
			Message:  fmt.Sprintf("failed to generate image with refs: %v", err),
		}
	}

	imageBytes, mimeType, err := extractGeneratedImage(resp)
	if err != nil {
		return nil, "", &SearchError{
			Provider: nanoBananaSource,
			Code:     "NO_RESULTS",
			Message:  err.Error(),
		}
	}

	return imageBytes, mimeType, nil
}

func extractGeneratedImage(response *genai.GenerateContentResponse) ([]byte, string, error) {
	if response == nil {
		return nil, "", fmt.Errorf("no response from Gemini")
	}

	for _, candidate := range response.Candidates {
		if candidate == nil || candidate.Content == nil {
			continue
		}

		for _, part := range candidate.Content.Parts {
			if part == nil || part.InlineData == nil || len(part.InlineData.Data) == 0 {
				continue
			}

			mimeType := part.InlineData.MIMEType
			if mimeType == "" {
				mimeType = "image/png"
			}

			return append([]byte(nil), part.InlineData.Data...), mimeType, nil
		}
	}

	return nil, "", fmt.Errorf("no image data returned from Gemini")
}

func encodeDataURL(imageBytes []byte, mimeType string) (string, error) {
	if len(imageBytes) == 0 {
		return "", fmt.Errorf("no image bytes returned")
	}

	normalizedBytes, err := normalizePNG(imageBytes, mimeType)
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("data:image/png;base64,%s", base64.StdEncoding.EncodeToString(normalizedBytes)), nil
}

func decodeDataURL(url string) (io.ReadCloser, error) {
	header, payload, ok := strings.Cut(url, ",")
	if !ok || !strings.HasPrefix(header, "data:") || !strings.Contains(header, ";base64") {
		return nil, fmt.Errorf("unsupported data URI: %s", url)
	}

	data, err := base64.StdEncoding.DecodeString(payload)
	if err != nil {
		return nil, fmt.Errorf("decode data URI: %w", err)
	}

	return io.NopCloser(bytes.NewReader(data)), nil
}

func normalizePNG(imageBytes []byte, mimeType string) ([]byte, error) {
	if strings.EqualFold(strings.TrimSpace(mimeType), "image/png") {
		return append([]byte(nil), imageBytes...), nil
	}

	img, _, err := image.Decode(bytes.NewReader(imageBytes))
	if err != nil {
		return nil, fmt.Errorf("decode generated image: %w", err)
	}

	var buffer bytes.Buffer
	if err := png.Encode(&buffer, img); err != nil {
		return nil, fmt.Errorf("encode generated image as png: %w", err)
	}

	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]
}

func (c *NanoBananaClient) modelName() string {
	if c == nil || c.config == nil || strings.TrimSpace(c.config.Model) == "" {
		return DefaultNanoBananaModel
	}

	return c.config.Model
}

func (c *NanoBananaClient) textModelName() string {
	if c == nil || c.config == nil || strings.TrimSpace(c.config.TextModel) == "" {
		return DefaultNanoBananaTextModel
	}

	return c.config.TextModel
}

func normalizeNanoBananaConfig(config *NanoBananaConfig) *NanoBananaConfig {
	normalized := &NanoBananaConfig{}
	if config != nil {
		*normalized = *config
	}

	normalized.APIKey = strings.TrimSpace(normalized.APIKey)
	normalized.Model = strings.TrimSpace(normalized.Model)
	normalized.TextModel = strings.TrimSpace(normalized.TextModel)

	if normalized.Model == "" {
		normalized.Model = DefaultNanoBananaModel
	}
	if normalized.TextModel == "" {
		normalized.TextModel = DefaultNanoBananaTextModel
	}

	return normalized
}