NextExpense/internal/ai/receipt.go

// Package ai provides receipt data extraction from images/PDFs using
// configurable AI providers (Gemini or OpenAI-compatible).
//
// Provider selection is done via environment variables:
//
//	AI_PROVIDER=gemini          (default, uses GEMINI_API_KEY)
//	AI_PROVIDER=openai          (uses OPENAI_API_KEY, AI_MODEL, AI_BASE_URL)
//	                           (also works with Ollama, LocalAI, etc.)
package ai

import (
	"fmt"
	"os"
	"strings"
)

// 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"`
}

// ValidCategories is the list of allowed expense categories the AI should
// classify receipts into.
var ValidCategories = []string{
	"Airfare",
	"Accommodation",
	"Meals Self",
	"Staff Meal",
	"Client Meal",
	"Travel - Taxi",
	"Travel - Phone",
	"Misc Travel",
	"Mobile / Office Phone",
	"Office Supplies",
	"Postage / Couriers",
	"Other Expenses",
	"Hotel",
	"Per Diem",
	"Visa Fees",
	"Connectivity (internet connections)",
}

// categoryPrompt returns the comma-separated category list for AI prompts.
func categoryPrompt() string {
	return `"Airfare", "Accommodation", "Meals Self", "Staff Meal", "Client Meal", "Travel - Taxi", "Travel - Phone", "Misc Travel", "Mobile / Office Phone", "Office Supplies", "Postage / Couriers", "Other Expenses", "Hotel", "Per Diem", "Visa Fees", "Connectivity (internet connections)"`
}

// Provider is the interface that wraps receipt extraction.
// Each provider (Gemini, OpenAI, Ollama) implements this interface.
type Provider interface {
	ExtractReceipt(imagePath string) (*ReceiptData, error)
}

// ExtractReceipt dispatches to the configured AI provider.
// The provider is selected based on the AI_PROVIDER environment variable.
func ExtractReceipt(imagePath string) (*ReceiptData, error) {
	provider := getProvider()
	return provider.ExtractReceipt(imagePath)
}

// getProvider returns the appropriate Provider based on environment config.
func getProvider() Provider {
	providerName := strings.ToLower(strings.TrimSpace(os.Getenv("AI_PROVIDER")))

	switch providerName {
	case "openai":
		return newOpenAIProvider()
	default:
		return newGeminiProvider()
	}
}

// envOrDefault returns the environment variable value or a default if unset.
func envOrDefault(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

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

// readFile reads the full contents of a file from disk.
func readFile(path string) ([]byte, error) {
	info, err := os.Stat(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, fmt.Errorf("receipt file not found: %w", err)
		}
		return nil, fmt.Errorf("stat file %q: %w", path, err)
	}
	if info.IsDir() {
		return nil, fmt.Errorf("readFile: %q is a directory, not a file", path)
	}
	if info.Size() > 10<<20 {
		return nil, fmt.Errorf("readFile: %q exceeds 10 MB limit", path)
	}
	return os.ReadFile(path)
}

// detectMimeType determines the MIME type from 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
	if 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 ""
}