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"
)

const (
	// DefaultNanoBananaModel is the Gemini image model used for Nano Banana generation.
	DefaultNanoBananaModel = "gemini-3.1-flash-image-preview"

	// DefaultNanoBananaTextModel is the Gemini text model used for translation and scene generation.
	DefaultNanoBananaTextModel = "gemini-2.5-flash"

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

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 {
	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) {
	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)
	}

	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 := nanoBananaGenerateImage(ctx, c, prompt)
	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 := 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 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
}

func (c *NanoBananaClient) resolveTranslation(ctx context.Context, opts *SearchOptions, translation string) (string, error) {
	if translation != "" {
		fmt.Printf("Using provided translation: %s -> %s\n", opts.Query, translation)
		return translation, nil
	}

	translation, err := c.translateBulgarianToEnglish(ctx, opts.Query)
	if err != nil {
		fmt.Printf("Translation failed: %v, using original word\n", err)
		return opts.Query, nil
	}

	return translation, nil
}

func (c *NanoBananaClient) resolvePrompt(ctx context.Context, opts *SearchOptions, translatedWord string) (string, error) {
	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
	}

	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) > 1000 {
			customPrompt = customPrompt[:997] + "..."
		}
		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
}

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

	selectedStyle := chooseArtisticStyle()
	if selectedStyle == defaultArtisticStyle {
		fmt.Printf("  No artistic styles available, using generic prompt\n")
	}
	fmt.Printf("  Using image style: %s\n", selectedStyle)

	var prompt string

	if scene != "" {
		fullPrompt := fmt.Sprintf(
			"Generate a %s educational flashcard image illustrating \"%s\". Scene: %s "+
				"The image should be educational and suitable for language learning flashcards. "+
				"Requirements: The main subject or concept 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 scene makes \"%s\" 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, subject, withTerminalPunctuation(scene), subject,
		)

		if len(fullPrompt) > maxImagePromptChars {
			prompt = fmt.Sprintf(
				"Generate a %s flashcard image illustrating \"%s\". Scene: %s "+
					"The image should be educational and suitable for language learning flashcards. "+
					"Requirements: The main subject or concept must be clearly visible, centered, well lit, and easy to identify.",
				selectedStyle, subject, withTerminalPunctuation(scene),
			)

			if len(prompt) > maxImagePromptChars {
				maxSceneLen := maxImagePromptChars - len(fmt.Sprintf(
					"Generate a %s flashcard image illustrating \"%s\". Scene:  "+
						"The image should be educational and suitable for language learning flashcards. "+
						"Requirements: The main subject or concept must be clearly visible, centered, well lit, and easy to identify.",
					selectedStyle, subject,
				))
				if maxSceneLen > 3 && len(scene) > maxSceneLen {
					scene = scene[:maxSceneLen] + "..."
				}
				prompt = fmt.Sprintf(
					"Generate a %s flashcard image illustrating \"%s\". Scene: %s "+
						"The image should be educational and suitable for language learning flashcards. "+
						"Requirements: The main subject or concept must be clearly visible, centered, well lit, and easy to identify.",
					selectedStyle, subject, withTerminalPunctuation(scene),
				)
			}
		} else {
			prompt = fullPrompt
		}
	} else {
		prompt = fmt.Sprintf(
			"Generate a %s educational flashcard image illustrating \"%s\". %s "+
				"The image should be educational and suitable for language learning flashcards. "+
				"Requirements: The main subject or concept must be clearly visible, easily recognizable, and prominent in the image. Show it prominently centered with excellent lighting and sharp focus. "+
				"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, subject, fallbackVisualDirection(subject),
		)
	}

	if len(prompt) > maxImagePromptChars {
		prompt = prompt[:997] + "..."
	}

	return prompt
}

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 := 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),
		0.3,
		50,
	)
	if err != nil {
		return "", fmt.Errorf("translation failed: %w", err)
	}

	fmt.Printf("Translated '%s' to '%s'\n", word, translation)
	return translation, nil
}

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 := 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 string) ([]byte, string, error) {
	cfg := &genai.GenerateContentConfig{
		ResponseModalities: []string{string(genai.ModalityImage)},
		ImageConfig: &genai.ImageConfig{
			AspectRatio: nanoBananaAspectRatio,
		},
	}

	resp, err := 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
}

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
}