ai-media-hub/backend/services/gemini.go

package services

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"mime"
	"net/http"
	"strings"
	"time"
)

type GeminiService struct {
	APIKey string
	Client *http.Client
}

type AIRecommendation struct {
	Title        string `json:"title"`
	Link         string `json:"link"`
	ThumbnailURL string `json:"thumbnailUrl"`
	Source       string `json:"source"`
	Reason       string `json:"reason"`
	Recommended  bool   `json:"recommended"`
}

func NewGeminiService(apiKey string) *GeminiService {
	return &GeminiService{
		APIKey: apiKey,
		Client: &http.Client{Timeout: 40 * time.Second},
	}
}

func (g *GeminiService) Recommend(query string, candidates []SearchResult) ([]AIRecommendation, error) {
	if g.APIKey == "" {
		return nil, fmt.Errorf("gemini api key is not configured")
	}
	if len(candidates) == 0 {
		return []AIRecommendation{}, nil
	}

	type geminiPart map[string]any
	parts := []geminiPart{
		{
			"text": `Analyze the provided images for the user's search intent. Return JSON only in this shape:
{"recommendations":[{"index":0,"reason":"short reason","recommended":true}]}
Mark only the best matches as recommended=true. Keep reasons concise. User query: ` + query,
		},
	}

	maxImages := min(len(candidates), 8)
	for idx := 0; idx < maxImages; idx++ {
		img, mimeType, err := fetchImageAsInlineData(g.Client, candidates[idx].ThumbnailURL)
		if err != nil {
			continue
		}
		parts = append(parts,
			geminiPart{"text": fmt.Sprintf("Candidate %d: title=%s source=%s link=%s", idx, candidates[idx].Title, candidates[idx].Source, candidates[idx].Link)},
			geminiPart{"inlineData": map[string]string{"mimeType": mimeType, "data": img}},
		)
	}

	body := map[string]any{
		"contents": []map[string]any{
			{"parts": parts},
		},
		"generationConfig": map[string]any{
			"responseMimeType": "application/json",
		},
	}

	rawBody, _ := json.Marshal(body)
	endpoint := "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key=" + g.APIKey
	resp, err := g.Client.Post(endpoint, "application/json", bytes.NewReader(rawBody))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 300 {
		data, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
		return nil, fmt.Errorf("gemini returned status %d: %s", resp.StatusCode, string(data))
	}

	var payload struct {
		Candidates []struct {
			Content struct {
				Parts []struct {
					Text string `json:"text"`
				} `json:"parts"`
			} `json:"content"`
		} `json:"candidates"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {
		return nil, err
	}
	if len(payload.Candidates) == 0 || len(payload.Candidates[0].Content.Parts) == 0 {
		return nil, fmt.Errorf("gemini returned no candidates")
	}

	var parsed struct {
		Recommendations []struct {
			Index       int    `json:"index"`
			Reason      string `json:"reason"`
			Recommended bool   `json:"recommended"`
		} `json:"recommendations"`
	}
	if err := json.Unmarshal([]byte(payload.Candidates[0].Content.Parts[0].Text), &parsed); err != nil {
		return nil, err
	}

	recommendations := make([]AIRecommendation, 0, len(parsed.Recommendations))
	for _, rec := range parsed.Recommendations {
		if rec.Index < 0 || rec.Index >= len(candidates) || !rec.Recommended {
			continue
		}
		src := candidates[rec.Index]
		recommendations = append(recommendations, AIRecommendation{
			Title:        src.Title,
			Link:         src.Link,
			ThumbnailURL: src.ThumbnailURL,
			Source:       src.Source,
			Reason:       rec.Reason,
			Recommended:  true,
		})
	}

	if len(recommendations) == 0 {
		for _, candidate := range candidates[:min(4, len(candidates))] {
			recommendations = append(recommendations, AIRecommendation{
				Title:        candidate.Title,
				Link:         candidate.Link,
				ThumbnailURL: candidate.ThumbnailURL,
				Source:       candidate.Source,
				Reason:       "Fallback result because Gemini returned no recommended items.",
				Recommended:  true,
			})
		}
	}

	return recommendations, nil
}

func fetchImageAsInlineData(client *http.Client, imageURL string) (string, string, error) {
	resp, err := client.Get(imageURL)
	if err != nil {
		return "", "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 300 {
		return "", "", fmt.Errorf("thumbnail fetch failed with %d", resp.StatusCode)
	}

	contentType := resp.Header.Get("Content-Type")
	mimeType, _, _ := mime.ParseMediaType(contentType)
	if mimeType == "" || !strings.HasPrefix(mimeType, "image/") {
		mimeType = "image/jpeg"
	}

	data, err := io.ReadAll(io.LimitReader(resp.Body, 2*1024*1024))
	if err != nil {
		return "", "", err
	}
	return base64.StdEncoding.EncodeToString(data), mimeType, nil
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}