NextExpense/internal/ai/deepseek.go

package ai

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"image"
	"image/jpeg"
	"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 (
	deepseekAPIURL = "https://api.deepseek.com/v1/chat/completions"
	deepseekModel  = "deepseek-v4-flash"
	requestTimeout = 120 * time.Second
	maxImageSize   = 300 // max dimension in pixels (width or height)
	jpegQuality    = 50  // JPEG compression quality (1-100)
)

// deepseekRequest matches the DeepSeek API request format.
type deepseekRequest struct {
	Model       string            `json:"model"`
	Messages    []deepseekMessage `json:"messages"`
	Temperature float64           `json:"temperature"`
}

type deepseekMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// deepseekResponse matches the DeepSeek API response format.
type deepseekResponse struct {
	Choices []deepseekChoice `json:"choices"`
}

type deepseekChoice struct {
	Message deepseekResponseMessage `json:"message"`
}

type deepseekResponseMessage struct {
	Content string `json:"content"`
}

// ExtractReceipt sends a receipt image to the DeepSeek API and parses
// the structured receipt data from the response.
//
// Since DeepSeek's chat API does not support the OpenAI-style image_url
// content type, the image is resized (max 800px) and JPEG-compressed,
// then base64-encoded and embedded directly in the text prompt.
func ExtractReceipt(imagePath string) (*ReceiptData, error) {
	// 1. Validate the image file exists and is readable.
	imageData, err := readImageFile(imagePath)
	if err != nil {
		return &ReceiptData{}, err
	}

	// 2. Get the API key from the environment.
	apiKey := os.Getenv("DEEPSEEK_API_KEY")
	if apiKey == "" {
		err := errors.New("DEEPSEEK_API_KEY environment variable is not set")
		log.Printf("ExtractReceipt: %v", err)
		return &ReceiptData{}, err
	}

	// 3. Decode, downscale, and re-encode the image to keep token cost manageable.
	compressed, err := compressImage(imageData)
	if err != nil {
		log.Printf("ExtractReceipt: image compression failed: %v", err)
		// Fall back to raw image if compression fails.
		compressed = imageData
	}

	// 4. Base64-encode the compressed image.
	b64 := base64.StdEncoding.EncodeToString(compressed)

	// 5. Build the text prompt with the base64 image data inline.
	prompt := fmt.Sprintf(
		`Analyze this receipt image (base64-encoded JPEG/PNG below). Extract the following fields as a strict JSON object with these exact keys:
- "amount": float (the total amount paid, e.g. 42.50)
- "currency": string (3-letter ISO currency code, e.g. EUR, USD, GBP)
- "merchant": string (the store or business name)
- "category": string (one of: Food, Travel, Lodging, Software, Other)
- "date": string (the receipt date in YYYY-MM-DD format)

Return ONLY valid JSON. No markdown, no explanation, no code fences.

Image data: %s`, b64)

	payload := deepseekRequest{
		Model:       deepseekModel,
		Temperature: 0.1,
		Messages: []deepseekMessage{
			{
				Role:    "user",
				Content: prompt,
			},
		},
	}

	body, err := json.Marshal(payload)
	if err != nil {
		log.Printf("ExtractReceipt: failed to marshal request: %v", err)
		return &ReceiptData{}, err
	}

	// 6. Send the POST request.
	req, err := http.NewRequest(http.MethodPost, deepseekAPIURL, bytes.NewReader(body))
	if err != nil {
		log.Printf("ExtractReceipt: failed to create request: %v", err)
		return &ReceiptData{}, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+apiKey)

	client := &http.Client{Timeout: requestTimeout}
	resp, err := client.Do(req)
	if err != nil {
		log.Printf("ExtractReceipt: API request failed: %v", err)
		return &ReceiptData{}, err
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		log.Printf("ExtractReceipt: failed to read response: %v", err)
		return &ReceiptData{}, err
	}

	if resp.StatusCode != http.StatusOK {
		err := fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(respBody))
		log.Printf("ExtractReceipt: %v", err)
		return &ReceiptData{}, err
	}

	// 7. Parse the DeepSeek response.
	var apiResp deepseekResponse
	if err := json.Unmarshal(respBody, &apiResp); err != nil {
		log.Printf("ExtractReceipt: failed to parse response JSON: %v", err)
		return &ReceiptData{}, err
	}

	if len(apiResp.Choices) == 0 {
		return &ReceiptData{}, errors.New("API response contains no choices")
	}

	contentStr := strings.TrimSpace(apiResp.Choices[0].Message.Content)

	// Strip markdown code fences if the model wrapped the JSON.
	contentStr = stripMarkdownFences(contentStr)

	// 8. Parse the JSON content into ReceiptData.
	var receipt ReceiptData
	if err := json.Unmarshal([]byte(contentStr), &receipt); err != nil {
		log.Printf("ExtractReceipt: failed to parse receipt JSON from content: %v", err)
		log.Printf("ExtractReceipt: raw content was: %s", contentStr)
		return &ReceiptData{}, err
	}

	log.Printf("ExtractReceipt: successfully extracted receipt data: merchant=%q amount=%.2f %s",
		receipt.Merchant, receipt.Amount, receipt.Currency)
	return &receipt, nil
}

// compressImage decodes the image, downscales it (preserving aspect ratio)
// to fit within maxImageSize, and re-encodes as JPEG with the configured quality.
func compressImage(data []byte) ([]byte, error) {
	// Detect format and decode.
	img, format, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("image decode: %w", err)
	}

	bounds := img.Bounds()
	w := bounds.Dx()
	h := bounds.Dy()

	// If the image is already small, just re-encode at reduced quality.
	if w <= maxImageSize && h <= maxImageSize {
		var buf bytes.Buffer
		if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: jpegQuality}); err != nil {
			return nil, fmt.Errorf("jpeg re-encode: %w", err)
		}
		// Only use if smaller than original.
		if buf.Len() < len(data) {
			return buf.Bytes(), nil
		}
		return data, nil
	}

	// Calculate new dimensions.
	ratio := float64(maxImageSize) / float64(max(w, h))
	newW := int(float64(w) * ratio)
	newH := int(float64(h) * ratio)
	if newW < 1 {
		newW = 1
	}
	if newH < 1 {
		newH = 1
	}

	// Downscale using simple bilinear-like averaging (via RGBA iteration).
	scaled := scaleImage(img, newW, newH)

	var buf bytes.Buffer
	if err := jpeg.Encode(&buf, scaled, &jpeg.Options{Quality: jpegQuality}); err != nil {
		return nil, fmt.Errorf("jpeg encode: %w", err)
	}

	log.Printf("ExtractReceipt: compressed image %dx%d %s -> %dx%d JPEG (%d bytes)",
		w, h, format, newW, newH, buf.Len())
	return buf.Bytes(), nil
}

// scaleImage performs a simple nearest-neighbour downscale of an image.
func scaleImage(src image.Image, newW, newH int) image.Image {
	bounds := src.Bounds()
	srcW := bounds.Dx()
	srcH := bounds.Dy()
	dst := image.NewRGBA(image.Rect(0, 0, newW, newH))

	for y := 0; y < newH; y++ {
		for x := 0; x < newW; x++ {
			sx := x * srcW / newW
			sy := y * srcH / newH
			dst.Set(x, y, src.At(sx, sy))
		}
	}
	return dst
}

// stripMarkdownFences removes markdown code fences (```json ... ```) from the
// model output, if present.
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, not an image: %s", path)
	}

	if info.Size() > 10<<20 {
		return nil, fmt.Errorf("image file 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
}