NextExpense/internal/ai/deepseek.go

package ai

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

// ReceiptData represents the structured data extracted from a receipt image.
type ReceiptData struct {
	Amount   float64 `json:"amount"`
	Currency string  `json:"currency"`
	Merchant string  `json:"merchant"`
	Category string  `json:"category"`
	Date     string  `json:"date"`
}

const (
	geminiAPIURL   = "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.1-flash-lite:generateContent"
	requestTimeout = 30 * time.Second
)

// --- Gemini API types ---

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      *geminiError      `json:"error,omitempty"`
}

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

type geminiResponseContent struct {
	Parts []geminiResponsePart `json:"parts"`
}

type geminiResponsePart struct {
	Text string `json:"text"`
}

type geminiError struct {
	Message string `json:"message"`
}

// ExtractReceipt sends a receipt image to Google Gemini Vision API and
// returns structured receipt data extracted from the image.
//
// Gemini supports native vision processing via inline_data, so the image
// is sent directly as base64-encoded data with its MIME type.
func ExtractReceipt(imagePath string) (*ReceiptData, error) {
	// 1. Read the image file.
	imageData, err := readImageFile(imagePath)
	if err != nil {
		return &ReceiptData{}, err
	}

	// 2. Detect MIME type from magic bytes.
	mimeType := detectMimeType(imageData)
	if mimeType == "" {
		mimeType = "image/jpeg" // safe default
	}

	// 3. Get the API key.
	apiKey := os.Getenv("GEMINI_API_KEY")
	if apiKey == "" {
		// Fall back to the key from .env.example.
		apiKey = "AQ.Ab8RN6IQjKTQofuKOW2TT5mZ0zwt8rFa8X3SHGyYyce4DrbBJw"
	}

	// 4. Build the Gemini request with the image as inline_data.
	b64Data := base64.StdEncoding.EncodeToString(imageData)

	payload := geminiRequest{
		Contents: []geminiContent{
			{
				Parts: []geminiPart{
					{
						Text: "Analyze this receipt image. Extract the following fields as a strict JSON object with these exact keys: \"merchant\" (string, store or business name), \"amount\" (number, total paid), \"currency\" (string, 3-letter code like KES, USD, EUR), \"category\" (string, one of: Food, Travel, Lodging, Software, Other), \"date\" (string, YYYY-MM-DD format). Return ONLY valid JSON. No markdown, no explanation, no code fences.",
					},
					{
						InlineData: &geminiFileData{
							MimeType: mimeType,
							Data:     b64Data,
						},
					},
				},
			},
		},
	}

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

	// 5. Send to Gemini API.
	req, err := http.NewRequest(http.MethodPost, geminiAPIURL, 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", apiKey)

	client := &http.Client{Timeout: requestTimeout}
	resp, err := client.Do(req)
	if err != nil {
		return &ReceiptData{}, fmt.Errorf("API request failed: %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 API returned status %d: %s", resp.StatusCode, string(respBody))
	}

	// 6. Parse the response.
	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 API error: %s", apiResp.Error.Message)
	}

	if len(apiResp.Candidates) == 0 {
		return &ReceiptData{}, errors.New("Gemini returned no candidates")
	}

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

	contentStr := strings.TrimSpace(parts[0].Text)
	contentStr = stripMarkdownFences(contentStr)

	// 7. Parse the JSON into ReceiptData.
	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: merchant=%q amount=%.2f %s category=%q date=%q",
		receipt.Merchant, receipt.Amount, receipt.Currency, receipt.Category, receipt.Date)
	return &receipt, nil
}

// detectMimeType determines the MIME type of an image from its magic bytes.
func detectMimeType(data []byte) string {
	if len(data) < 4 {
		return ""
	}
	// JPEG
	if data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF {
		return "image/jpeg"
	}
	// PNG
	if data[0] == 0x89 && data[1] == 0x50 && data[2] == 0x4E && data[3] == 0x47 {
		return "image/png"
	}
	// WebP
	if len(data) >= 12 && data[0] == 0x52 && data[1] == 0x49 && data[2] == 0x46 &&
		data[3] == 0x46 && data[8] == 0x57 && data[9] == 0x45 && data[10] == 0x42 && data[11] == 0x50 {
		return "image/webp"
	}
	// GIF
	if data[0] == 0x47 && data[1] == 0x49 && data[2] == 0x46 {
		return "image/gif"
	}
	// BMP
	if data[0] == 0x42 && data[1] == 0x4D {
		return "image/bmp"
	}
	// TIFF
	if (data[0] == 0x49 && data[1] == 0x49 && data[2] == 0x2A && data[3] == 0x00) ||
		(data[0] == 0x4D && data[1] == 0x4D && data[2] == 0x00 && data[3] == 0x2A) {
		return "image/tiff"
	}
	// PDF: 25 50 44 46 (%PDF)
	if len(data) >= 4 && data[0] == 0x25 && data[1] == 0x50 && data[2] == 0x44 && data[3] == 0x46 {
		return "application/pdf"
	}

	// HEIC/HEIF (ftyp box at offset 4)
	if len(data) >= 12 && data[4] == 0x66 && data[5] == 0x74 && data[6] == 0x79 && data[7] == 0x70 {
		brand := string(data[8:12])
		switch brand {
		case "heic", "heix", "hevc", "hevx", "mif1", "msf1":
			return "image/heic"
		case "avif":
			return "image/avif"
		}
	}
	return ""
}

// stripMarkdownFences removes markdown code fences from model output.
func stripMarkdownFences(s string) string {
	s = strings.TrimSpace(s)
	if strings.HasPrefix(s, "```") {
		s = s[3:]
		if idx := strings.Index(s, "\n"); idx != -1 {
			s = s[idx+1:]
		}
	}
	if strings.HasSuffix(s, "```") {
		s = s[:len(s)-3]
	}
	return strings.TrimSpace(s)
}

// readImageFile reads the full contents of an image file from disk.
func readImageFile(path string) ([]byte, error) {
	info, err := os.Stat(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, fmt.Errorf("image file does not exist: %s", path)
		}
		return nil, fmt.Errorf("cannot stat image file %s: %w", path, err)
	}
	if info.IsDir() {
		return nil, fmt.Errorf("path is a directory: %s", path)
	}
	if info.Size() > 10<<20 {
		return nil, fmt.Errorf("image too large: %d bytes (max 10MB)", info.Size())
	}
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("cannot read image file %s: %w", path, err)
	}
	return data, nil
}