MoneyMap/MoneyMap/Services/AIReceiptParser.cs

using Microsoft.EntityFrameworkCore;
using MoneyMap.Data;
using MoneyMap.Models;
using System.Text.Json;

namespace MoneyMap.Services
{
    public interface IReceiptParser
    {
        Task<ReceiptParseResult> ParseReceiptAsync(long receiptId, string? model = null, string? notes = null);
    }

    public class AIReceiptParser : IReceiptParser
    {
        private readonly MoneyMapContext _db;
        private readonly IReceiptManager _receiptManager;
        private readonly IPdfToImageConverter _pdfConverter;
        private readonly IAIVisionClientResolver _clientResolver;
        private readonly IMerchantService _merchantService;
        private readonly IServiceProvider _serviceProvider;
        private readonly IConfiguration _configuration;
        private readonly ILogger<AIReceiptParser> _logger;
        private string? _promptTemplate;

        public AIReceiptParser(
            MoneyMapContext db,
            IReceiptManager receiptManager,
            IPdfToImageConverter pdfConverter,
            IAIVisionClientResolver clientResolver,
            IMerchantService merchantService,
            IServiceProvider serviceProvider,
            IConfiguration configuration,
            ILogger<AIReceiptParser> logger)
        {
            _db = db;
            _receiptManager = receiptManager;
            _pdfConverter = pdfConverter;
            _clientResolver = clientResolver;
            _merchantService = merchantService;
            _serviceProvider = serviceProvider;
            _configuration = configuration;
            _logger = logger;
        }

        public async Task<ReceiptParseResult> ParseReceiptAsync(long receiptId, string? model = null, string? notes = null)
        {
            var receipt = await _db.Receipts
                .Include(r => r.Transaction)
                .FirstOrDefaultAsync(r => r.Id == receiptId);

            if (receipt == null)
                return ReceiptParseResult.Failure("Receipt not found.");

            var filePath = _receiptManager.GetReceiptPhysicalPath(receipt);
            if (!File.Exists(filePath))
                return ReceiptParseResult.Failure("Receipt file not found on disk.");

            var selectedModel = model ?? _configuration["AI:ReceiptParsingModel"] ?? "gpt-4o-mini";
            var (client, provider) = _clientResolver.Resolve(selectedModel);

            var parseLog = new ReceiptParseLog
            {
                ReceiptId = receiptId,
                Provider = provider,
                Model = selectedModel,
                StartedAtUtc = DateTime.UtcNow,
                Success = false
            };

            try
            {
                var (base64Data, mediaType) = await PrepareImageDataAsync(receipt, filePath);
                var promptText = await BuildPromptAsync(receipt, notes);
                var visionResult = await client.AnalyzeImageAsync(base64Data, mediaType, promptText, selectedModel);

                if (!visionResult.IsSuccess)
                {
                    await SaveParseLogAsync(parseLog, visionResult.ErrorMessage);
                    return ReceiptParseResult.Failure(visionResult.ErrorMessage!);
                }

                var parseData = ParseResponse(visionResult.Content);
                await ApplyParseResultAsync(receipt, receiptId, parseData, notes);

                parseLog.Success = true;
                parseLog.Confidence = parseData.Confidence;
                parseLog.RawProviderPayloadJson = JsonSerializer.Serialize(parseData);
                await SaveParseLogAsync(parseLog);

                await TryAutoMapReceiptAsync(receipt, receiptId);

                var lineCount = parseData.LineItems.Count;
                return ReceiptParseResult.Success($"Parsed {lineCount} line items from receipt.");
            }
            catch (Exception ex)
            {
                await SaveParseLogAsync(parseLog, ex.Message);
                _logger.LogError(ex, "Error parsing receipt {ReceiptId}: {Message}", receiptId, ex.Message);
                return ReceiptParseResult.Failure($"Error parsing receipt: {ex.Message}");
            }
        }

        private async Task<(string Base64Data, string MediaType)> PrepareImageDataAsync(Receipt receipt, string filePath)
        {
            if (receipt.ContentType == "application/pdf")
            {
                var base64 = await _pdfConverter.ConvertFirstPageToBase64Async(filePath);
                return (base64, "image/png");
            }

            var fileBytes = await File.ReadAllBytesAsync(filePath);
            return (Convert.ToBase64String(fileBytes), receipt.ContentType);
        }

        private async Task<string> BuildPromptAsync(Receipt receipt, string? userNotes = null)
        {
            var promptText = await LoadPromptTemplateAsync();

            var transactionName = receipt.Transaction?.Name;
            if (!string.IsNullOrWhiteSpace(transactionName))
            {
                promptText += $"\n\nNote: This transaction was recorded as \"{transactionName}\" in the bank statement, which may help identify the merchant if the receipt is unclear.";
            }

            var parsingNotes = _configuration["AI:ReceiptParsingNotes"];
            if (!string.IsNullOrWhiteSpace(parsingNotes))
            {
                promptText += $"\n\nAdditional notes: {parsingNotes}";
            }

            if (!string.IsNullOrWhiteSpace(userNotes))
            {
                promptText += $"\n\nUser notes for this receipt: {userNotes}";
            }

            promptText += "\n\nRespond ONLY with valid JSON, no other text.";
            return promptText;
        }

        private static ParsedReceiptData ParseResponse(string? content)
        {
            if (string.IsNullOrWhiteSpace(content))
                return new ParsedReceiptData();

            return JsonSerializer.Deserialize<ParsedReceiptData>(content, new JsonSerializerOptions
            {
                PropertyNameCaseInsensitive = true
            }) ?? new ParsedReceiptData();
        }

        private async Task ApplyParseResultAsync(Receipt receipt, long receiptId, ParsedReceiptData parseData, string? notes)
        {
            // Update receipt fields
            receipt.ParsingNotes = notes;
            receipt.Merchant = parseData.Merchant;
            receipt.Total = parseData.Total;
            receipt.Subtotal = parseData.Subtotal;
            receipt.Tax = parseData.Tax;
            receipt.ReceiptDate = parseData.ReceiptDate;
            receipt.DueDate = parseData.DueDate;

            // Update transaction merchant if needed
            if (receipt.Transaction != null &&
                !string.IsNullOrWhiteSpace(parseData.Merchant) &&
                receipt.Transaction.MerchantId == null)
            {
                var merchantId = await _merchantService.GetOrCreateIdAsync(parseData.Merchant);
                receipt.Transaction.MerchantId = merchantId;
            }

            // Replace line items
            var existingItems = await _db.ReceiptLineItems
                .Where(li => li.ReceiptId == receiptId)
                .ToListAsync();
            _db.ReceiptLineItems.RemoveRange(existingItems);

            var lineItems = parseData.LineItems.Select((item, index) => new ReceiptLineItem
            {
                ReceiptId = receiptId,
                LineNumber = index + 1,
                Description = item.Description,
                Sku = item.Upc,
                Quantity = item.Quantity,
                UnitPrice = item.UnitPrice,
                LineTotal = item.LineTotal,
                Voided = item.Voided
            }).ToList();

            _db.ReceiptLineItems.AddRange(lineItems);
            await _db.SaveChangesAsync();
        }

        private async Task SaveParseLogAsync(ReceiptParseLog parseLog, string? error = null)
        {
            parseLog.Error = error;
            parseLog.CompletedAtUtc = DateTime.UtcNow;
            _db.ReceiptParseLogs.Add(parseLog);
            await _db.SaveChangesAsync();
        }

        private async Task TryAutoMapReceiptAsync(Receipt receipt, long receiptId)
        {
            if (receipt.TransactionId.HasValue)
                return;

            try
            {
                using var scope = _serviceProvider.CreateScope();
                var autoMapper = scope.ServiceProvider.GetRequiredService<IReceiptAutoMapper>();
                await autoMapper.AutoMapReceiptAsync(receiptId);
                _logger.LogInformation("Auto-mapping completed for receipt {ReceiptId}", receiptId);
            }
            catch (Exception ex)
            {
                _logger.LogWarning(ex, "Auto-mapping failed for receipt {ReceiptId}: {Message}", receiptId, ex.Message);
            }
        }

        private async Task<string> LoadPromptTemplateAsync()
        {
            if (_promptTemplate != null)
                return _promptTemplate;

            var promptPath = Path.Combine(AppContext.BaseDirectory, "Prompts", "ReceiptParserPrompt.txt");

            if (!File.Exists(promptPath))
                throw new FileNotFoundException($"Receipt parser prompt template not found at: {promptPath}");

            _promptTemplate = await File.ReadAllTextAsync(promptPath);
            return _promptTemplate;
        }
    }

    /// <summary>
    /// Resolves the appropriate AI vision client based on model name.
    /// </summary>
    public interface IAIVisionClientResolver
    {
        (IAIVisionClient Client, string Provider) Resolve(string model);
    }

    public class AIVisionClientResolver : IAIVisionClientResolver
    {
        private readonly OpenAIVisionClient _openAIClient;
        private readonly ClaudeVisionClient _claudeClient;
        private readonly OllamaVisionClient _ollamaClient;
        private readonly LlamaCppVisionClient _llamaCppClient;

        public AIVisionClientResolver(
            OpenAIVisionClient openAIClient,
            ClaudeVisionClient claudeClient,
            OllamaVisionClient ollamaClient,
            LlamaCppVisionClient llamaCppClient)
        {
            _openAIClient = openAIClient;
            _claudeClient = claudeClient;
            _ollamaClient = ollamaClient;
            _llamaCppClient = llamaCppClient;
        }

        public (IAIVisionClient Client, string Provider) Resolve(string model)
        {
            if (model.StartsWith("llamacpp:"))
                return (_llamaCppClient, "LlamaCpp");

            if (model.StartsWith("ollama:"))
                return (_ollamaClient, "Ollama");

            if (model.StartsWith("claude-"))
                return (_claudeClient, "Anthropic");

            return (_openAIClient, "OpenAI");
        }
    }

    public class ParsedReceiptData
    {
        public string? Merchant { get; set; }
        public DateTime? ReceiptDate { get; set; }
        public DateTime? DueDate { get; set; }
        public decimal? Subtotal { get; set; }
        public decimal? Tax { get; set; }
        public decimal? Total { get; set; }
        public decimal Confidence { get; set; } = 0.5m;
        public List<ParsedLineItem> LineItems { get; set; } = new();
    }

    public class ParsedLineItem
    {
        public string Description { get; set; } = "";
        public string? Upc { get; set; }
        public decimal? Quantity { get; set; }
        public decimal? UnitPrice { get; set; }
        public decimal LineTotal { get; set; }
        public bool Voided { get; set; }
    }

    public class ReceiptParseResult
    {
        public bool IsSuccess { get; init; }
        public string? Message { get; init; }

        public static ReceiptParseResult Success(string message) =>
            new() { IsSuccess = true, Message = message };

        public static ReceiptParseResult Failure(string message) =>
            new() { IsSuccess = false, Message = message };
    }
}