NextExpense/internal/ai/deepseek.go

package ai

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"image"
	"image/jpeg"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"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.
// If Go's built-in decoders cannot handle the format (e.g. HEIC from iPhones),
// it attempts to convert via external tools (heif-convert or ImageMagick).
func compressImage(data []byte) ([]byte, error) {
	// Try Go's built-in image decoders first (covers JPEG, PNG, GIF, BMP, TIFF, WebP).
	img, format, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		// Built-in decoder failed — try external tools for HEIC and other formats.
		log.Printf("compressImage: Go decoder failed (%v), trying external conversion", err)
		converted, err := convertWithExternalTool(data)
		if err != nil {
			return nil, fmt.Errorf("image decode and external conversion both failed: %w", err)
		}
		// Try decoding the converted data.
		img, format, err = image.Decode(bytes.NewReader(converted))
		if err != nil {
			return nil, fmt.Errorf("decoding converted image also failed: %w", err)
		}
		data = converted
	}

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

// convertWithExternalTool tries to convert an unsupported image format (e.g. HEIC)
// to JPEG using heif-convert or ImageMagick's convert command.
func convertWithExternalTool(data []byte) ([]byte, error) {
	// Write the unknown image to a temp file (needed for CLI tools).
	tmpDir := os.TempDir()
	inPath := filepath.Join(tmpDir, "ef-convert-in-"+fmt.Sprintf("%d", time.Now().UnixNano()))
	outPath := filepath.Join(tmpDir, "ef-convert-out-"+fmt.Sprintf("%d", time.Now().UnixNano())+".jpg")

	if err := os.WriteFile(inPath, data, 0644); err != nil {
		return nil, fmt.Errorf("write temp input: %w", err)
	}
	defer os.Remove(inPath)
	defer os.Remove(outPath)

	// Try heif-convert first (fastest, handles HEIC natively).
	if heifErr := tryHEIFConvert(inPath, outPath); heifErr == nil {
		outData, err := os.ReadFile(outPath)
		if err == nil && len(outData) > 0 {
			log.Printf("convertWithExternalTool: heif-convert succeeded (%d bytes)", len(outData))
			return outData, nil
		}
	}

	// Fall back to ImageMagick convert.
	if magickErr := exec.Command("convert", inPath, "-resize", fmt.Sprintf("%dx%d>", maxImageSize, maxImageSize), "-quality", fmt.Sprintf("%d", jpegQuality), outPath).Run(); magickErr == nil {
		outData, err := os.ReadFile(outPath)
		if err == nil && len(outData) > 0 {
			log.Printf("convertWithExternalTool: ImageMagick convert succeeded (%d bytes)", len(outData))
			return outData, nil
		}
	}

	return nil, errors.New("all external conversion tools failed")
}

// tryHEIFConvert attempts to convert a HEIC/HEIF file to JPEG using heif-convert.
func tryHEIFConvert(inPath, outPath string) error {
	// heif-convert outputs to {input}.jpg by default.
	defaultOut := inPath + ".jpg"
	defer os.Remove(defaultOut)

	cmd := exec.Command("heif-convert", inPath, outPath)
	if output, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("heif-convert failed: %w, output: %s", err, string(output))
	}

	// If it wrote to the default path instead of our outPath, move it.
	if _, err := os.Stat(outPath); os.IsNotExist(err) {
		if _, err := os.Stat(defaultOut); err == nil {
			os.Rename(defaultOut, outPath)
		}
	}

	_, err := os.Stat(outPath)
	return err
}

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