NextExpense/internal/ai/deepseek.go

package ai

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"image"
	"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 = 60 * time.Second
)

// 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 extracts structured receipt data from an image using a two-step
// pipeline: OCR (Tesseract) → LLM parsing (DeepSeek).
//
// Step 1 — OCR: Run Tesseract on the receipt image to extract raw text.
// Step 2 — LLM: Send the OCR text to DeepSeek to parse into structured JSON.
//
// This works with any image format (JPEG, PNG, HEIC, etc.) and produces
// real extracted data rather than AI-hallucinated values.
func ExtractReceipt(imagePath string) (*ReceiptData, error) {
	apiKey := os.Getenv("DEEPSEEK_API_KEY")
	if apiKey == "" {
		return &ReceiptData{}, errors.New("DEEPSEEK_API_KEY environment variable is not set")
	}

	// Step 1: OCR — extract text from the receipt image.
	ocrText, err := ocrImage(imagePath)
	if err != nil {
		log.Printf("ExtractReceipt: OCR failed: %v", err)
		return &ReceiptData{}, fmt.Errorf("OCR failed: %w", err)
	}

	ocrText = strings.TrimSpace(ocrText)
	if ocrText == "" {
		return &ReceiptData{}, errors.New("OCR returned no text from the receipt image")
	}

	log.Printf("ExtractReceipt: OCR extracted %d characters of text", len(ocrText))

	// Step 2: LLM — send OCR text to DeepSeek for structured parsing.
	receipt, err := parseReceiptText(ocrText, apiKey)
	if err != nil {
		log.Printf("ExtractReceipt: LLM parsing failed: %v", err)
		return &ReceiptData{}, err
	}

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

// ocrImage runs Tesseract OCR on the image file and returns the extracted text.
func ocrImage(imagePath string) (string, error) {
	// First, try to convert the image to a format Tesseract handles well.
	// For HEIC files, this also handles the conversion.
	jpegPath, err := ensureJPEG(imagePath)
	if err != nil {
		return "", fmt.Errorf("image preparation: %w", err)
	}
	if jpegPath != imagePath {
		defer os.Remove(jpegPath)
	}

	cmd := exec.Command("tesseract", jpegPath, "stdout", "--psm", "6")
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		return "", fmt.Errorf("tesseract failed: %w, stderr: %s", err, strings.TrimSpace(stderr.String()))
	}

	return stdout.String(), nil
}

// ensureJPEG converts the image at the given path to JPEG if it isn't already.
// Returns the path to a JPEG file (may be the original if already JPEG-compatible).
// The caller should remove the returned path if it differs from the input path.
func ensureJPEG(path string) (string, error) {
	// Read the file to check format.
	data, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}

	// Try to decode with Go's stdlib — if it works, we can re-encode as JPEG.
	// Tesseract can handle PNG/BMP/TIFF natively, so only convert if Go
	// cannot decode the format (e.g., HEIC).
	_, _, err = image.Decode(bytes.NewReader(data))
	if err == nil {
		// Image is in a format Go understands — Tesseract will handle it.
		// No conversion needed unless it's a very large image.
		return path, nil
	}

	// Go couldn't decode it (likely HEIC). Try external conversion.
	log.Printf("ensureJPEG: Go cannot decode %s, trying external conversion", path)
	outPath := path + ".ocr-convert.jpg"
	if err := convertToJPEG(data, outPath); err != nil {
		return "", fmt.Errorf("external conversion failed: %w", err)
	}

	return outPath, nil
}

// convertToJPEG tries multiple tools to convert image data to JPEG.
func convertToJPEG(data []byte, outPath string) error {
	tmpDir := os.TempDir()
	inPath := filepath.Join(tmpDir, "ef-convert-"+fmt.Sprintf("%d", time.Now().UnixNano()))
	if err := os.WriteFile(inPath, data, 0644); err != nil {
		return err
	}
	defer os.Remove(inPath)

	// Try heif-convert first.
	if err := exec.Command("heif-convert", inPath, outPath).Run(); err == nil {
		if _, err := os.Stat(outPath); err == nil {
			return nil
		}
	}

	// Fall back to ImageMagick.
	if err := exec.Command("convert", inPath, outPath).Run(); err == nil {
		if _, err := os.Stat(outPath); err == nil {
			return nil
		}
	}

	return errors.New("all image conversion tools failed (tried heif-convert, convert)")
}

// parseReceiptText sends raw OCR text to DeepSeek and returns structured data.
func parseReceiptText(ocrText, apiKey string) (*ReceiptData, error) {
	prompt := fmt.Sprintf(`Extract receipt information from the following OCR text.

Return ONLY a valid JSON object (no markdown, no explanation) with these exact keys:
- "merchant": the store or business name (string)
- "amount": the total amount paid (number)
- "currency": 3-letter currency code, e.g. KES, USD, EUR (string)
- "category": one of: Food, Travel, Lodging, Software, Other (string)
- "date": the receipt date in YYYY-MM-DD format (string)

If a field is not found in the text, use null for the value.
Do NOT make up or invent data that is not present in the OCR text.

OCR text from receipt:
---
%s
---`, ocrText)

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

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

	req, err := http.NewRequest(http.MethodPost, deepseekAPIURL, bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("create request: %w", 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 {
		return nil, fmt.Errorf("API request failed: %w", err)
	}
	defer resp.Body.Close()

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

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(respBody))
	}

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

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

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

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

	return &receipt, nil
}

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