MoneyMap/MoneyMap/Pages/AICategorizePreview.cshtml.cs

using Microsoft.AspNetCore.Mvc;
using Microsoft.AspNetCore.Mvc.RazorPages;
using Microsoft.EntityFrameworkCore;
using MoneyMap.Data;
using MoneyMap.Models;
using MoneyMap.Services;
using System.Text.Json;

namespace MoneyMap.Pages
{
    public class AICategorizePreviewModel : PageModel
    {
        private readonly MoneyMapContext _db;
        private readonly ITransactionAICategorizer _aiCategorizer;
        private readonly IConfiguration _config;

        public AICategorizePreviewModel(
            MoneyMapContext db,
            ITransactionAICategorizer aiCategorizer,
            IConfiguration config)
        {
            _db = db;
            _aiCategorizer = aiCategorizer;
            _config = config;
        }

        public List<ProposalViewModel> Proposals { get; set; } = new();
        public string ModelUsed { get; set; } = "";
        public string AIProvider
        {
            get
            {
                var model = SelectedModel;
                if (model.StartsWith("llamacpp:", StringComparison.OrdinalIgnoreCase)) return "LlamaCpp";
                if (model.StartsWith("ollama:", StringComparison.OrdinalIgnoreCase)) return "Ollama";
                if (model.StartsWith("claude-", StringComparison.OrdinalIgnoreCase)) return "Anthropic";
                return "OpenAI";
            }
        }
        public string SelectedModel => _config["AI:ReceiptParsingModel"] ?? "gpt-4o-mini";

        [TempData]
        public string? StoredTransactionIds { get; set; }

        public long[]? TransactionIds { get; set; }

        [TempData]
        public string? ProposalsJson { get; set; }

        [TempData]
        public string? SuccessMessage { get; set; }

        [TempData]
        public string? ErrorMessage { get; set; }

        public int SelectedTransactionCount { get; set; }

        public async Task<IActionResult> OnGetAsync()
        {
            // Load transaction IDs from TempData if available
            if (!string.IsNullOrEmpty(StoredTransactionIds))
            {
                TransactionIds = JsonSerializer.Deserialize<long[]>(StoredTransactionIds);
                // Keep in TempData for the generate form
                StoredTransactionIds = JsonSerializer.Serialize(TransactionIds);
            }

            // Count selected transactions if IDs provided
            if (TransactionIds != null && TransactionIds.Length > 0)
            {
                SelectedTransactionCount = await _db.Transactions
                    .CountAsync(t => TransactionIds.Contains(t.Id));
            }

            // Check if we have proposals from TempData
            if (!string.IsNullOrEmpty(ProposalsJson))
            {
                var storedProposals = JsonSerializer.Deserialize<List<StoredProposal>>(ProposalsJson);
                if (storedProposals != null)
                {
                    await LoadProposalViewModels(storedProposals);
                }
            }

            return Page();
        }

        public IActionResult OnPostStoreIds(long[] transactionIds)
        {
            if (transactionIds == null || transactionIds.Length == 0)
            {
                return RedirectToPage("/Transactions");
            }

            StoredTransactionIds = JsonSerializer.Serialize(transactionIds);
            return RedirectToPage();
        }

        public async Task<IActionResult> OnPostGenerateAsync()
        {
            var uncategorized = await _db.Transactions
                .Include(t => t.Card)
                .Include(t => t.Account)
                .Include(t => t.Merchant)
                .Include(t => t.TransferToAccount)
                .Where(t => string.IsNullOrWhiteSpace(t.Category))
                .OrderByDescending(t => t.Date)
                .Take(50)
                .ToListAsync();

            if (uncategorized.Count == 0)
            {
                ErrorMessage = "No uncategorized transactions to process.";
                return RedirectToPage("/Recategorize");
            }

            var proposals = await _aiCategorizer.ProposeBatchCategorizationAsync(uncategorized, SelectedModel);

            if (proposals.Count == 0)
            {
                ErrorMessage = "AI could not generate any categorization proposals.";
                return RedirectToPage("/Recategorize");
            }

            // Store proposals in TempData
            var storedProposals = proposals.Select(p => new StoredProposal
            {
                TransactionId = p.TransactionId,
                Category = p.Category,
                CanonicalMerchant = p.CanonicalMerchant,
                Pattern = p.Pattern,
                Priority = p.Priority,
                Confidence = p.Confidence,
                Reasoning = p.Reasoning,
                CreateRule = p.CreateRule
            }).ToList();

            ProposalsJson = JsonSerializer.Serialize(storedProposals);
            ModelUsed = SelectedModel;

            return RedirectToPage();
        }

        public async Task<IActionResult> OnPostGenerateForIdsAsync(long[]? transactionIds)
        {
            // Try to get IDs from form first, then from TempData
            if ((transactionIds == null || transactionIds.Length == 0) && !string.IsNullOrEmpty(StoredTransactionIds))
            {
                transactionIds = JsonSerializer.Deserialize<long[]>(StoredTransactionIds);
            }

            if (transactionIds == null || transactionIds.Length == 0)
            {
                ErrorMessage = "No transactions selected.";
                return RedirectToPage("/Transactions");
            }

            var transactions = await _db.Transactions
                .Include(t => t.Card)
                .Include(t => t.Account)
                .Include(t => t.Merchant)
                .Include(t => t.TransferToAccount)
                .Where(t => transactionIds.Contains(t.Id))
                .OrderByDescending(t => t.Date)
                .Take(50)
                .ToListAsync();

            if (transactions.Count == 0)
            {
                ErrorMessage = "Selected transactions not found.";
                return RedirectToPage("/Transactions");
            }

            var proposals = await _aiCategorizer.ProposeBatchCategorizationAsync(transactions, SelectedModel);

            if (proposals.Count == 0)
            {
                ErrorMessage = "AI could not generate any categorization proposals.";
                return RedirectToPage("/Transactions");
            }

            // Store proposals in TempData
            var storedProposals = proposals.Select(p => new StoredProposal
            {
                TransactionId = p.TransactionId,
                Category = p.Category,
                CanonicalMerchant = p.CanonicalMerchant,
                Pattern = p.Pattern,
                Priority = p.Priority,
                Confidence = p.Confidence,
                Reasoning = p.Reasoning,
                CreateRule = p.CreateRule
            }).ToList();

            ProposalsJson = JsonSerializer.Serialize(storedProposals);
            ModelUsed = SelectedModel;

            return RedirectToPage();
        }

        public async Task<IActionResult> OnPostApplyAsync(long[] selectedIds, long[] createRules, string? proposalsData)
        {
            // Read proposals from the hidden form field (not TempData, which can be lost on app restart)
            var json = proposalsData ?? ProposalsJson;

            if (string.IsNullOrEmpty(json))
            {
                ErrorMessage = "No proposals to apply. Please generate new suggestions.";
                return RedirectToPage("/Recategorize");
            }

            var storedProposals = JsonSerializer.Deserialize<List<StoredProposal>>(json);
            if (storedProposals == null || storedProposals.Count == 0)
            {
                ErrorMessage = "No proposals to apply (deserialization returned empty).";
                return RedirectToPage("/Recategorize");
            }

            var selectedSet = selectedIds?.ToHashSet() ?? new HashSet<long>();
            var createRulesSet = createRules?.ToHashSet() ?? new HashSet<long>();

            if (selectedSet.Count == 0)
            {
                ErrorMessage = $"No transactions were selected. ({storedProposals.Count} proposals available but 0 selectedIds received from form)";
                return RedirectToPage("/Recategorize");
            }

            int applied = 0;
            int rulesCreated = 0;
            int rulesUpdated = 0;
            var errors = new List<string>();

            foreach (var stored in storedProposals)
            {
                if (!selectedSet.Contains(stored.TransactionId))
                    continue;

                var proposal = new AICategoryProposal
                {
                    TransactionId = stored.TransactionId,
                    Category = stored.Category,
                    CanonicalMerchant = stored.CanonicalMerchant,
                    Pattern = stored.Pattern,
                    Priority = stored.Priority,
                    Confidence = stored.Confidence,
                    Reasoning = stored.Reasoning,
                    CreateRule = stored.CreateRule
                };

                // Check if user wants to create/update rule for this one
                var shouldCreateRule = createRulesSet.Contains(stored.TransactionId);

                var result = await _aiCategorizer.ApplyProposalAsync(stored.TransactionId, proposal, shouldCreateRule);
                if (result.Success)
                {
                    applied++;
                    if (result.RuleCreated)
                        rulesCreated++;
                    if (result.RuleUpdated)
                        rulesUpdated++;
                }
                else
                {
                    errors.Add($"ID {stored.TransactionId}: {result.ErrorMessage}");
                }
            }

            var parts = new List<string> { $"Applied {applied} of {selectedSet.Count} selected categorizations" };
            if (rulesCreated > 0) parts.Add($"created {rulesCreated} new rules");
            if (rulesUpdated > 0) parts.Add($"updated {rulesUpdated} existing rules");
            SuccessMessage = string.Join(". ", parts) + ".";

            if (errors.Count > 0)
            {
                ErrorMessage = "Some proposals failed: " + string.Join("; ", errors);
            }

            return RedirectToPage("/Recategorize");
        }

        private async Task LoadProposalViewModels(List<StoredProposal> storedProposals)
        {
            var transactionIds = storedProposals.Select(p => p.TransactionId).ToList();
            var transactions = await _db.Transactions
                .Where(t => transactionIds.Contains(t.Id))
                .ToDictionaryAsync(t => t.Id);

            // Look up existing rules for all proposed patterns
            var proposedPatterns = storedProposals
                .Where(p => !string.IsNullOrWhiteSpace(p.Pattern))
                .Select(p => p.Pattern!)
                .Distinct()
                .ToList();

            var existingRules = await _db.CategoryMappings
                .Where(m => proposedPatterns.Contains(m.Pattern))
                .ToDictionaryAsync(m => m.Pattern, m => m.Category);

            foreach (var stored in storedProposals)
            {
                if (transactions.TryGetValue(stored.TransactionId, out var txn))
                {
                    string? existingCategory = null;
                    var hasExisting = !string.IsNullOrWhiteSpace(stored.Pattern)
                        && existingRules.TryGetValue(stored.Pattern!, out existingCategory);

                    Proposals.Add(new ProposalViewModel
                    {
                        TransactionId = stored.TransactionId,
                        TransactionName = txn.Name,
                        TransactionMemo = txn.Memo,
                        TransactionAmount = txn.Amount,
                        TransactionDate = txn.Date,
                        CurrentCategory = txn.Category,
                        ProposedCategory = stored.Category,
                        ProposedMerchant = stored.CanonicalMerchant,
                        ProposedPattern = stored.Pattern,
                        Confidence = stored.Confidence,
                        Reasoning = stored.Reasoning,
                        CreateRule = stored.CreateRule,
                        HasExistingRule = hasExisting,
                        ExistingRuleCategory = existingCategory
                    });
                }
            }

            // Keep proposals in TempData for the apply action
            ProposalsJson = JsonSerializer.Serialize(storedProposals);
        }

        public class ProposalViewModel
        {
            public long TransactionId { get; set; }
            public string TransactionName { get; set; } = "";
            public string? TransactionMemo { get; set; }
            public decimal TransactionAmount { get; set; }
            public DateTime TransactionDate { get; set; }
            public string? CurrentCategory { get; set; }
            public string ProposedCategory { get; set; } = "";
            public string? ProposedMerchant { get; set; }
            public string? ProposedPattern { get; set; }
            public decimal Confidence { get; set; }
            public string? Reasoning { get; set; }
            public bool CreateRule { get; set; }
            public bool HasExistingRule { get; set; }
            public string? ExistingRuleCategory { get; set; }

            /// <summary>
            /// True when the pattern exists but is mapped to a different category than proposed.
            /// </summary>
            public bool NeedsRuleUpdate => HasExistingRule && ExistingRuleCategory != ProposedCategory;
        }

        public class StoredProposal
        {
            public long TransactionId { get; set; }
            public string Category { get; set; } = "";
            public string? CanonicalMerchant { get; set; }
            public string? Pattern { get; set; }
            public int Priority { get; set; }
            public decimal Confidence { get; set; }
            public string? Reasoning { get; set; }
            public bool CreateRule { get; set; }
        }
    }
}