NextExpense/internal/ai/gemini.go

package ai

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"time"
)

const geminiTimeout = 30 * time.Second

type geminiRequest struct {
	Contents []geminiContent `json:"contents"`
}

type geminiContent struct {
	Parts []geminiPart `json:"parts"`
}

type geminiPart struct {
	Text       string          `json:"text,omitempty"`
	InlineData *geminiFileData `json:"inline_data,omitempty"`
}

type geminiFileData struct {
	MimeType string `json:"mime_type"`
	Data     string `json:"data"`
}

type geminiResponse struct {
	Candidates []geminiCandidate `json:"candidates"`
	Error      *struct {
		Message string `json:"message"`
	} `json:"error,omitempty"`
}

type geminiCandidate struct {
	Content geminiResponseContent `json:"content"`
}

type geminiResponseContent struct {
	Parts []struct {
		Text string `json:"text"`
	} `json:"parts"`
}

type geminiProvider struct {
	apiKey string
	apiURL string
}

func newGeminiProvider() geminiProvider {
	apiKey := os.Getenv("GEMINI_API_KEY")
	model := os.Getenv("GEMINI_MODEL")
	if model == "" {
		model = "gemini-3.1-flash-lite"
	}
	return geminiProvider{
		apiKey: apiKey,
		apiURL: fmt.Sprintf("https://generativelanguage.googleapis.com/v1beta/models/%s:generateContent", model),
	}
}

func (p geminiProvider) ExtractReceipt(imagePath string) (*ReceiptData, error) {
	if p.apiKey == "" {
		return &ReceiptData{}, errors.New("GEMINI_API_KEY environment variable not set")
	}

	imageData, err := readFile(imagePath)
	if err != nil {
		return &ReceiptData{}, fmt.Errorf("read file: %w", err)
	}

	mimeType := detectMimeType(imageData)
	if mimeType == "" {
		mimeType = "image/jpeg"
	}

	b64Data := base64.StdEncoding.EncodeToString(imageData)

	payload := geminiRequest{
		Contents: []geminiContent{{
			Parts: []geminiPart{
				{Text: fmt.Sprintf("Analyze this receipt. Extract as strict JSON with keys: \"merchant\" (string), \"amount\" (number), \"currency\" (3-letter code), \"category\" (string, MUST be exactly one of: %s), \"date\" (YYYY-MM-DD). Return ONLY valid JSON. No markdown.", categoryPrompt())},
				{InlineData: &geminiFileData{MimeType: mimeType, Data: b64Data}},
			},
		}},
	}

	body, err := json.Marshal(payload)
	if err != nil {
		return &ReceiptData{}, fmt.Errorf("marshal request: %w", err)
	}

	req, err := http.NewRequest(http.MethodPost, p.apiURL, bytes.NewReader(body))
	if err != nil {
		return &ReceiptData{}, fmt.Errorf("create request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-goog-api-key", p.apiKey)

	client := &http.Client{Timeout: geminiTimeout}
	resp, err := client.Do(req)
	if err != nil {
		return &ReceiptData{}, fmt.Errorf("API request: %w", err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return &ReceiptData{}, fmt.Errorf("read response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return &ReceiptData{}, fmt.Errorf("Gemini status %d: %s", resp.StatusCode, strings.TrimSpace(string(respBody)))
	}

	var apiResp geminiResponse
	if err := json.Unmarshal(respBody, &apiResp); err != nil {
		return &ReceiptData{}, fmt.Errorf("parse response: %w", err)
	}

	if apiResp.Error != nil {
		return &ReceiptData{}, fmt.Errorf("Gemini error: %s", apiResp.Error.Message)
	}
	if len(apiResp.Candidates) == 0 {
		return &ReceiptData{}, errors.New("no candidates in Gemini response")
	}

	parts := apiResp.Candidates[0].Content.Parts
	if len(parts) == 0 {
		return &ReceiptData{}, errors.New("no response text from Gemini")
	}

	contentStr := stripMarkdownFences(parts[0].Text)
	var receipt ReceiptData
	if err := json.Unmarshal([]byte(contentStr), &receipt); err != nil {
		return &ReceiptData{}, fmt.Errorf("parse receipt JSON: %w (content: %s)", err, contentStr)
	}

	log.Printf("ExtractReceipt [gemini]: merchant=%q amount=%.2f %s category=%q date=%q",
		receipt.Merchant, receipt.Amount, receipt.Currency, receipt.Category, receipt.Date)
	return &receipt, nil
}