path: root/internal/image/openai.go

package image

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/sashabaranov/go-openai"
)

// OpenAIClient implements ImageSearcher for OpenAI DALL-E image generation
type OpenAIClient struct {
	client     *openai.Client
	apiKey     string
	model      string // dall-e-2 or dall-e-3
	size       string // 256x256, 512x512, 1024x1024
	quality    string // standard or hd (dall-e-3 only)
	style      string // natural or vivid (dall-e-3 only)
	lastPrompt string // Store the last used prompt for attribution

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

// OpenAIConfig holds configuration for the OpenAI image provider
type OpenAIConfig struct {
	APIKey  string
	Model   string
	Size    string
	Quality string
	Style   string
}

// NewOpenAIClient creates a new OpenAI DALL-E client
func NewOpenAIClient(config *OpenAIConfig) *OpenAIClient {
	if config.APIKey == "" {
		// Return nil client that will fail on operations
		return &OpenAIClient{}
	}

	client := openai.NewClient(config.APIKey)

	// Set defaults
	if config.Model == "" {
		config.Model = "dall-e-3"
	}
	if config.Size == "" {
		config.Size = "1024x1024"
	}
	if config.Quality == "" {
		config.Quality = "standard"
	}
	if config.Style == "" {
		config.Style = "natural"
	}

	oc := &OpenAIClient{
		client:  client,
		apiKey:  config.APIKey,
		model:   config.Model,
		size:    config.Size,
		quality: config.Quality,
		style:   config.Style,
	}

	return oc
}

// Search generates an image for the Bulgarian word using DALL-E
func (c *OpenAIClient) Search(ctx context.Context, opts *SearchOptions) ([]SearchResult, error) {
	if c.client == nil {
		return nil, &SearchError{
			Provider: "openai",
			Code:     "NO_API_KEY",
			Message:  "OpenAI API key not configured",
		}
	}

	// Use provided translation if available, otherwise translate Bulgarian word to English
	var translatedWord string
	if opts.Translation != "" {
		// Use the translation that was already provided (from UI or user input)
		translatedWord = opts.Translation
		fmt.Printf("Using provided translation: %s -> %s\n", opts.Query, translatedWord)
	} else {
		// Translate Bulgarian word to English for better results
		var err error
		translatedWord, err = c.translateBulgarianToEnglish(ctx, opts.Query)
		if err != nil {
			// If translation fails, fall back to using the original word
			fmt.Printf("Translation failed: %v, using original word\n", err)
			translatedWord = opts.Query
		}
	}

	// Create prompt - use custom if provided, otherwise generate educational prompt
	var prompt string
	if opts.CustomPrompt != "" && strings.TrimSpace(opts.CustomPrompt) != "" {
		prompt = strings.TrimSpace(opts.CustomPrompt)
		// Ensure custom prompt doesn't exceed 1000 characters
		if len(prompt) > 1000 {
			prompt = prompt[:997] + "..."
			fmt.Printf("Custom prompt truncated to 1000 chars\n")
		}
		fmt.Printf("Using custom prompt: %s\n", prompt)
	} else {
		prompt = c.createEducationalPrompt(ctx, opts.Query, translatedWord)
		if prompt == "" {
			return nil, &SearchError{
				Provider: "openai",
				Code:     "PROMPT_GENERATION_FAILED",
				Message:  "Failed to generate image prompt - artistic styles could not be loaded",
			}
		}
	}

	// Store the prompt for attribution
	c.lastPrompt = prompt

	// Call the callback if set
	if c.PromptCallback != nil {
		c.PromptCallback(prompt)
	}

	// Log the prompt to stdout for debugging
	fmt.Printf("OpenAI Image Generation Prompt (%d chars): %s\n", len(prompt), prompt)
	fmt.Printf("OpenAI Image Generation: Using model '%s' with size '%s'\n", c.model, c.size)

	// Create the image generation request
	req := openai.ImageRequest{
		Prompt:         prompt,
		Model:          c.model,
		Size:           c.size,
		ResponseFormat: openai.CreateImageResponseFormatURL,
		N:              1,
	}

	// Add model-specific parameters
	if c.model == "dall-e-3" {
		req.Quality = c.quality
		req.Style = c.style
	}

	// Generate the image
	resp, err := c.client.CreateImage(ctx, req)
	if err != nil {
		return nil, &SearchError{
			Provider: "openai",
			Code:     "API_ERROR",
			Message:  fmt.Sprintf("Failed to generate image: %v", err),
		}
	}

	if len(resp.Data) == 0 {
		return nil, &SearchError{
			Provider: "openai",
			Code:     "NO_RESULTS",
			Message:  "No image generated",
		}
	}

	// Get the generated image URL
	imageURL := resp.Data[0].URL

	// Create result
	result := SearchResult{
		ID:           c.generateImageID(opts.Query),
		URL:          imageURL,
		ThumbnailURL: imageURL,
		Width:        c.getSizeWidth(),
		Height:       c.getSizeHeight(),
		Description:  fmt.Sprintf("Generated educational image for %s (%s)", opts.Query, translatedWord),
		Attribution:  "Generated by OpenAI DALL-E",
		Source:       "openai",
	}

	return []SearchResult{result}, nil
}

// Download downloads an image from the given URL
func (c *OpenAIClient) Download(ctx context.Context, url string) (io.ReadCloser, error) {
	// Download from URL
	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err != nil {
		return nil, err
	}

	resp, err := http.DefaultClient.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 the required attribution text
func (c *OpenAIClient) GetAttribution(result *SearchResult) string {
	var attribution strings.Builder
	attribution.WriteString("Image generated by OpenAI DALL-E\n\n")
	fmt.Fprintf(&attribution, "Model: %s\n", c.model)
	fmt.Fprintf(&attribution, "Size: %s\n", c.size)
	if c.model == "dall-e-3" {
		fmt.Fprintf(&attribution, "Quality: %s\n", c.quality)
		fmt.Fprintf(&attribution, "Style: %s\n", c.style)
	}
	fmt.Fprintf(&attribution, "\nPrompt used:\n%s\n", c.lastPrompt)
	fmt.Fprintf(&attribution, "\nGenerated at: %s\n", time.Now().Format("2006-01-02 15:04:05"))
	return attribution.String()
}

// Name returns the name of the provider
func (c *OpenAIClient) Name() string {
	return "openai"
}

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

// SetPromptCallback sets a callback function that will be called when the prompt is generated
func (c *OpenAIClient) SetPromptCallback(callback func(prompt string)) {
	c.PromptCallback = callback
}

// createEducationalPrompt generates a prompt optimized for language learning
func (c *OpenAIClient) createEducationalPrompt(ctx context.Context, bulgarianWord, englishTranslation string) string {
	// Generate a scene description for the word
	scene, err := c.generateSceneDescription(ctx, bulgarianWord, englishTranslation)
	if err != nil {
		fmt.Printf("  Failed to generate scene: %v, using basic prompt\n", err)
		scene = ""
	}

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

	// Define prompt components in order of importance
	var prompt string

	if scene != "" {
		// Full prompt with scene
		fullPrompt := fmt.Sprintf(
			"Generate a %s depicting: %s. "+
				"The image should be educational and suitable for language learning flashcards. "+
				"Requirements: The main subject must be clearly visible, easily recognizable, and prominent in the image. It should occupy the central area with sharp focus and proper lighting. Ensure the subject is shown from an angle that makes it immediately identifiable. "+
				"IMPORTANT: No text whatsoever. Do not include any words, letters, typography, labels, captions, or writing of any kind. Image only, without any text elements.",
			selectedStyle, scene,
		)

		// Check if full prompt exceeds 1000 characters
		if len(fullPrompt) > 1000 {
			// Try without the IMPORTANT notice
			prompt = fmt.Sprintf(
				"Generate a %s depicting: %s. "+
					"The image should be educational and suitable for language learning flashcards. "+
					"Requirements: The main subject must be clearly visible, easily recognizable, and prominent in the image. It should occupy the central area with sharp focus and proper lighting.",
				selectedStyle, scene,
			)

			// If still too long, truncate the scene
			if len(prompt) > 1000 {
				// Truncate scene to fit within limit
				maxSceneLen := 1000 - len(fmt.Sprintf(
					"Generate a %s depicting: . "+
						"The image should be educational and suitable for language learning flashcards. "+
						"Requirements: The main subject must be clearly visible, easily recognizable, and prominent in the image.",
					selectedStyle,
				))
				if len(scene) > maxSceneLen {
					scene = scene[:maxSceneLen] + "..."
				}
				prompt = fmt.Sprintf(
					"Generate a %s depicting: %s. "+
						"The image should be educational and suitable for language learning flashcards. "+
						"Requirements: The main subject must be clearly visible, easily recognizable, and prominent in the image.",
					selectedStyle, scene,
				)
			}
		} else {
			prompt = fullPrompt
		}
	} else {
		// Basic prompt without scene
		prompt = fmt.Sprintf(
			"Generate a %s of %s. "+
				"The image should be educational and suitable for language learning flashcards. "+
				"Requirements: The %s must be clearly visible and easily recognizable. Show it prominently centered with excellent lighting and sharp focus.",
			selectedStyle, englishTranslation, englishTranslation,
		)
	}

	// Final check to ensure prompt is within 1000 characters
	if len(prompt) > 1000 {
		prompt = prompt[:997] + "..."
	}

	return prompt
}

// translateBulgarianToEnglish translates a Bulgarian word to English using OpenAI
func (c *OpenAIClient) translateBulgarianToEnglish(ctx context.Context, word string) (string, error) {
	// Use OpenAI chat completion to translate
	fmt.Printf("OpenAI Translation: Using model 'gpt-4o-mini' to translate '%s'\n", word)

	req := openai.ChatCompletionRequest{
		Model: openai.GPT4oMini,
		Messages: []openai.ChatCompletionMessage{
			{
				Role:    openai.ChatMessageRoleUser,
				Content: fmt.Sprintf("Translate the Bulgarian word '%s' to English. Respond with only the English translation, nothing else.", word),
			},
		},
		Temperature: 0.3, // Lower temperature for more consistent translations
		MaxTokens:   50,
	}

	resp, err := c.client.CreateChatCompletion(ctx, req)
	if err != nil {
		return "", fmt.Errorf("translation failed: %w", err)
	}

	if len(resp.Choices) == 0 || resp.Choices[0].Message.Content == "" {
		return "", fmt.Errorf("no translation received")
	}

	translation := strings.TrimSpace(resp.Choices[0].Message.Content)
	fmt.Printf("Translated '%s' to '%s'\n", word, translation)

	return translation, nil
}

// generateSceneDescription generates a contextual scene description for the word
func (c *OpenAIClient) generateSceneDescription(ctx context.Context, bulgarianWord, englishTranslation string) (string, error) {
	// Use OpenAI to generate a scene description
	fmt.Printf("OpenAI Scene Generation: Creating scene for '%s' (%s)\n", bulgarianWord, englishTranslation)

	req := openai.ChatCompletionRequest{
		Model: openai.GPT4oMini,
		Messages: []openai.ChatCompletionMessage{
			{
				Role:    openai.ChatMessageRoleSystem,
				Content: "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.",
			},
			{
				Role:    openai.ChatMessageRoleUser,
				Content: 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),
			},
		},
		Temperature: 0.7, // Balanced temperature for creativity with consistency
		MaxTokens:   100,
	}

	resp, err := c.client.CreateChatCompletion(ctx, req)
	if err != nil {
		return "", fmt.Errorf("scene generation failed: %w", err)
	}

	if len(resp.Choices) == 0 || resp.Choices[0].Message.Content == "" {
		return "", fmt.Errorf("no scene description received")
	}

	scene := strings.TrimSpace(resp.Choices[0].Message.Content)
	fmt.Printf("Generated scene: %s\n", scene)

	return scene, nil
}

// generateImageID generates a unique ID for the image
func (c *OpenAIClient) generateImageID(word string) string {
	// Create hash of the word for unique ID
	hash := md5.Sum([]byte(word))
	return hex.EncodeToString(hash[:])[:8]
}

// getSizeWidth returns the width based on the configured size
func (c *OpenAIClient) getSizeWidth() int {
	switch c.size {
	case "256x256":
		return 256
	case "512x512":
		return 512
	case "1024x1024":
		return 1024
	default:
		return 1024
	}
}

// getSizeHeight returns the height based on the configured size
func (c *OpenAIClient) getSizeHeight() int {
	// All DALL-E sizes are square
	return c.getSizeWidth()
}