Purrse/src/Purrse.Api/Services/AiCategorizationService.cs

using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization;
using Microsoft.EntityFrameworkCore;
using Purrse.Core.DTOs;
using Purrse.Core.Enums;
using Purrse.Core.Interfaces.Services;
using Purrse.Core.Models;
using Purrse.Data;

namespace Purrse.Api.Services;

public class AiCategorizationService : IAiCategorizationService
{
    private readonly PurrseDbContext _db;
    private readonly IHttpClientFactory _httpClientFactory;
    private readonly ILogger<AiCategorizationService> _logger;

    internal const string DefaultSystemPrompt =
        """
        You are a personal finance transaction categorizer. Classify each transaction into the most appropriate category.

        You MUST respond with valid JSON in exactly this format:
        {
          "classifications": [
            {"index": 0, "categoryName": "Groceries", "confidence": 0.9},
            {"index": 1, "categoryName": "Utilities", "confidence": 0.85}
          ]
        }

        Rules:
        - "index" must be the transaction number (integer) from the list below
        - "categoryName" must be one of the exact category names listed below
        - "confidence" must be a number between 0 and 1
        - Include one entry for every transaction

        Available categories:
        {categories}
        """;

    internal const string DefaultUserPromptTemplate =
        """
        {examples}Classify these uncategorized transactions:
        {transactions}
        """;

    private static readonly JsonSerializerOptions JsonOptions = new()
    {
        PropertyNameCaseInsensitive = true,
        DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull
    };

    public AiCategorizationService(
        PurrseDbContext db,
        IHttpClientFactory httpClientFactory,
        ILogger<AiCategorizationService> logger)
    {
        _db = db;
        _httpClientFactory = httpClientFactory;
        _logger = logger;
    }

    public async Task<AiCategorizationSettingsResponse> GetSettingsAsync(Guid userId)
    {
        var settings = await _db.AiCategorizationSettings.FirstOrDefaultAsync(s => s.UserId == userId);

        if (settings == null)
        {
            return new AiCategorizationSettingsResponse(
                false, "http://localhost:11434", "llama3.1:8b", 0.7, 30, true, null, null, 16384, null);
        }

        return new AiCategorizationSettingsResponse(
            settings.IsEnabled, settings.OllamaUrl, settings.ModelName,
            settings.ConfidenceThreshold, settings.TimeoutSeconds, settings.UpdatePayeeDefaults,
            settings.SystemPrompt, settings.UserPromptTemplate,
            settings.ChatContextSize, settings.ChatBotName);
    }

    public async Task<AiCategorizationSettingsResponse> UpdateSettingsAsync(Guid userId, UpdateAiCategorizationSettingsRequest request)
    {
        var settings = await _db.AiCategorizationSettings.FirstOrDefaultAsync(s => s.UserId == userId);

        if (settings == null)
        {
            settings = new AiCategorizationSettings
            {
                Id = Guid.NewGuid(),
                UserId = userId
            };
            _db.AiCategorizationSettings.Add(settings);
        }

        settings.IsEnabled = request.IsEnabled;
        settings.OllamaUrl = request.OllamaUrl;
        settings.ModelName = request.ModelName;
        settings.ConfidenceThreshold = request.ConfidenceThreshold;
        settings.TimeoutSeconds = request.TimeoutSeconds;
        settings.UpdatePayeeDefaults = request.UpdatePayeeDefaults;
        settings.SystemPrompt = string.IsNullOrWhiteSpace(request.SystemPrompt) ? null : request.SystemPrompt;
        settings.UserPromptTemplate = string.IsNullOrWhiteSpace(request.UserPromptTemplate) ? null : request.UserPromptTemplate;
        settings.ChatContextSize = request.ChatContextSize;
        settings.ChatBotName = string.IsNullOrWhiteSpace(request.ChatBotName) ? null : request.ChatBotName;
        settings.UpdatedAt = DateTime.UtcNow;

        await _db.SaveChangesAsync();

        return new AiCategorizationSettingsResponse(
            settings.IsEnabled, settings.OllamaUrl, settings.ModelName,
            settings.ConfidenceThreshold, settings.TimeoutSeconds, settings.UpdatePayeeDefaults,
            settings.SystemPrompt, settings.UserPromptTemplate,
            settings.ChatContextSize, settings.ChatBotName);
    }

    public async Task<OllamaConnectionTestResponse> TestConnectionAsync(Guid userId)
    {
        var settings = await _db.AiCategorizationSettings.FirstOrDefaultAsync(s => s.UserId == userId);
        var ollamaUrl = settings?.OllamaUrl ?? "http://localhost:11434";

        try
        {
            var client = _httpClientFactory.CreateClient("Ollama");
            client.Timeout = TimeSpan.FromSeconds(10);

            var response = await client.GetAsync($"{ollamaUrl.TrimEnd('/')}/api/tags");
            response.EnsureSuccessStatusCode();

            var json = await response.Content.ReadAsStringAsync();
            var tagsResponse = JsonSerializer.Deserialize<OllamaTagsResponse>(json, JsonOptions);

            var modelNames = tagsResponse?.Models?.Select(m => m.Name).ToList() ?? new List<string>();

            return new OllamaConnectionTestResponse(true, null, modelNames);
        }
        catch (Exception ex)
        {
            _logger.LogWarning(ex, "Ollama connection test failed for user {UserId}", userId);
            return new OllamaConnectionTestResponse(false, ex.Message, null);
        }
    }

    public async Task<List<TransactionClassification>> ClassifyTransactionsAsync(Guid userId, List<Transaction> uncategorizedTransactions)
    {
        if (uncategorizedTransactions.Count == 0)
            return new List<TransactionClassification>();

        try
        {
            var settings = await _db.AiCategorizationSettings.FirstOrDefaultAsync(s => s.UserId == userId);
            if (settings == null || !settings.IsEnabled)
                return new List<TransactionClassification>();

            // Load all active categories with parent names
            var categories = await _db.Categories
                .Include(c => c.Parent)
                .Where(c => c.UserId == userId && c.IsActive)
                .ToListAsync();

            if (categories.Count == 0)
                return new List<TransactionClassification>();

            // Load 25 recent categorized transactions as few-shot examples
            var recentCategorized = await _db.Transactions
                .Include(t => t.Category).ThenInclude(c => c!.Parent)
                .Include(t => t.Account)
                .Where(t => t.Account.UserId == userId && t.CategoryId != null)
                .OrderByDescending(t => t.Date)
                .Take(25)
                .ToListAsync();

            var allClassifications = new List<TransactionClassification>();

            // Process in batches of 50
            for (int batchStart = 0; batchStart < uncategorizedTransactions.Count; batchStart += 50)
            {
                var batch = uncategorizedTransactions.Skip(batchStart).Take(50).ToList();
                var batchClassifications = await ClassifyBatchAsync(settings, categories, recentCategorized, batch, batchStart);
                allClassifications.AddRange(batchClassifications);
            }

            return allClassifications;
        }
        catch (Exception ex)
        {
            _logger.LogWarning(ex, "AI categorization failed for user {UserId}", userId);
            return new List<TransactionClassification>();
        }
    }

    private async Task<List<TransactionClassification>> ClassifyBatchAsync(
        AiCategorizationSettings settings,
        List<Category> categories,
        List<Transaction> recentCategorized,
        List<Transaction> batch,
        int globalOffset)
    {
        var systemMessage = BuildSystemMessage(settings, categories);
        var userMessage = BuildUserMessage(settings, recentCategorized, batch);

        var chatRequest = new OllamaChatRequest(
            settings.ModelName,
            new List<OllamaChatMessage>
            {
                new("system", systemMessage),
                new("user", userMessage)
            },
            false,
            "json");

        try
        {
            var client = _httpClientFactory.CreateClient("Ollama");
            client.Timeout = TimeSpan.FromSeconds(settings.TimeoutSeconds);

            var requestJson = JsonSerializer.Serialize(chatRequest, JsonOptions);
            var content = new StringContent(requestJson, Encoding.UTF8, "application/json");

            _logger.LogInformation("Sending {Count} transactions to Ollama for classification (batch offset {Offset})", batch.Count, globalOffset);

            var response = await client.PostAsync($"{settings.OllamaUrl.TrimEnd('/')}/api/chat", content);
            response.EnsureSuccessStatusCode();

            var responseJson = await response.Content.ReadAsStringAsync();
            var chatResponse = JsonSerializer.Deserialize<OllamaChatResponse>(responseJson, JsonOptions);

            if (chatResponse?.Message?.Content == null)
            {
                _logger.LogWarning("Ollama returned null message content for batch at offset {Offset}", globalOffset);
                return new List<TransactionClassification>();
            }

            _logger.LogDebug("Ollama raw response for batch at offset {Offset}: {Response}", globalOffset, chatResponse.Message.Content);

            return ParseClassifications(chatResponse.Message.Content, categories, batch.Count, globalOffset, settings.ConfidenceThreshold);
        }
        catch (Exception ex)
        {
            _logger.LogWarning(ex, "Ollama API call failed for batch starting at {Offset}", globalOffset);
            return new List<TransactionClassification>();
        }
    }

    private static string BuildSystemMessage(AiCategorizationSettings settings, List<Category> categories)
    {
        var prompt = settings.SystemPrompt ?? DefaultSystemPrompt;

        var sb = new StringBuilder();
        foreach (var category in categories)
        {
            var parentInfo = category.Parent != null ? $" (under \"{category.Parent.Name}\")" : "";
            sb.AppendLine($"- \"{category.Name}\"{parentInfo}");
        }
        var categoryList = sb.ToString();

        return prompt.Contains("{categories}")
            ? prompt.Replace("{categories}", categoryList)
            : prompt + "\n" + categoryList;
    }

    private static string BuildUserMessage(AiCategorizationSettings settings, List<Transaction> recentCategorized, List<Transaction> batch)
    {
        var prompt = settings.UserPromptTemplate ?? DefaultUserPromptTemplate;

        // Build examples block
        var examplesSb = new StringBuilder();
        if (recentCategorized.Count > 0)
        {
            examplesSb.AppendLine("Recent categorized transactions (for reference):");
            for (int i = 0; i < recentCategorized.Count; i++)
            {
                var t = recentCategorized[i];
                var categoryName = t.Category?.Parent != null
                    ? $"{t.Category.Parent.Name} > {t.Category.Name}"
                    : t.Category?.Name ?? "Unknown";
                examplesSb.AppendLine($"{i + 1}. \"{t.PayeeName}\" (${t.Amount:F2}) -> \"{categoryName}\"");
            }
            examplesSb.AppendLine();
        }
        var examples = examplesSb.ToString();

        // Build transactions block
        var txnSb = new StringBuilder();
        for (int i = 0; i < batch.Count; i++)
        {
            var t = batch[i];
            var type = t.Amount >= 0 ? "credit" : "debit";
            txnSb.AppendLine($"{i}. \"{t.PayeeName}\" (${t.Amount:F2}, {type}, {t.Date:yyyy-MM-dd}, memo: \"{t.Memo}\")");
        }
        var transactions = txnSb.ToString();

        var result = prompt;
        if (result.Contains("{examples}"))
            result = result.Replace("{examples}", examples);
        else
            result = examples + result;

        if (result.Contains("{transactions}"))
            result = result.Replace("{transactions}", transactions);
        else
            result = result + "\n" + transactions;

        return result;
    }

    private List<TransactionClassification> ParseClassifications(
        string responseContent,
        List<Category> categories,
        int batchSize,
        int globalOffset,
        double confidenceThreshold)
    {
        try
        {
            using var doc = JsonDocument.Parse(responseContent);
            var root = doc.RootElement;

            // Try "classifications" wrapper first, then bare array
            JsonElement classificationsArray;
            if (root.TryGetProperty("classifications", out classificationsArray))
            {
                // good
            }
            else if (root.ValueKind == JsonValueKind.Array)
            {
                classificationsArray = root;
            }
            else
            {
                // Try common alternative wrapper names
                if (!root.TryGetProperty("results", out classificationsArray) &&
                    !root.TryGetProperty("data", out classificationsArray))
                {
                    _logger.LogWarning("Ollama response has no 'classifications' array. Root properties: {Props}",
                        string.Join(", ", root.EnumerateObject().Select(p => p.Name)));
                    return new List<TransactionClassification>();
                }
            }

            var categoryLookup = categories.ToDictionary(c => c.Id);
            var categoryNameLookup = categories
                .GroupBy(c => c.Name.Trim().ToLowerInvariant())
                .ToDictionary(g => g.Key, g => g.First());

            var results = new List<TransactionClassification>();
            int skippedNoIndex = 0, skippedBadIndex = 0, skippedLowConf = 0, skippedNoMatch = 0;

            foreach (var item in classificationsArray.EnumerateArray())
            {
                // Parse index — handle both number and string
                int localIndex;
                if (item.TryGetProperty("index", out var indexProp))
                {
                    if (indexProp.ValueKind == JsonValueKind.Number)
                        localIndex = indexProp.GetInt32();
                    else if (indexProp.ValueKind == JsonValueKind.String && int.TryParse(indexProp.GetString(), out var parsed))
                        localIndex = parsed;
                    else { skippedNoIndex++; continue; }
                }
                else { skippedNoIndex++; continue; }

                if (localIndex < 0 || localIndex >= batchSize)
                { skippedBadIndex++; continue; }

                var globalIndex = globalOffset + localIndex;

                // Parse confidence — handle number, string, or missing (default 1.0)
                double confidence = 1.0;
                if (item.TryGetProperty("confidence", out var confProp))
                {
                    if (confProp.ValueKind == JsonValueKind.Number)
                        confidence = confProp.GetDouble();
                    else if (confProp.ValueKind == JsonValueKind.String && double.TryParse(confProp.GetString(), out var parsedConf))
                        confidence = parsedConf;
                }

                if (confidence < confidenceThreshold)
                { skippedLowConf++; continue; }

                Guid? categoryId = null;
                string? categoryName = null;

                // Try to resolve by categoryId (also check snake_case variant)
                foreach (var propName in new[] { "categoryId", "category_id" })
                {
                    if (categoryId.HasValue) break;
                    if (item.TryGetProperty(propName, out var catIdProp) && catIdProp.ValueKind == JsonValueKind.String)
                    {
                        if (Guid.TryParse(catIdProp.GetString(), out var parsedId) && categoryLookup.ContainsKey(parsedId))
                        {
                            categoryId = parsedId;
                            categoryName = categoryLookup[parsedId].Name;
                        }
                    }
                }

                // Fall back to name matching (also check snake_case variant)
                if (!categoryId.HasValue)
                {
                    string? rawName = null;
                    foreach (var propName in new[] { "categoryName", "category_name", "category" })
                    {
                        if (item.TryGetProperty(propName, out var catNameProp) && catNameProp.ValueKind == JsonValueKind.String)
                        {
                            rawName = catNameProp.GetString();
                            break;
                        }
                    }

                    if (rawName != null)
                    {
                        var name = rawName.Trim().ToLowerInvariant();

                        // Exact match
                        if (categoryNameLookup.TryGetValue(name, out var matched))
                        {
                            categoryId = matched.Id;
                            categoryName = matched.Name;
                        }
                        else
                        {
                            // Handle "Parent > Child" format — try the last segment
                            var lastSegment = name.Contains('>') ? name.Split('>').Last().Trim() : null;
                            if (lastSegment != null && categoryNameLookup.TryGetValue(lastSegment, out var segmentMatch))
                            {
                                categoryId = segmentMatch.Id;
                                categoryName = segmentMatch.Name;
                            }
                            else
                            {
                                // Contains-based fuzzy match as last resort
                                var fuzzy = categories.FirstOrDefault(c =>
                                    c.Name.Trim().ToLowerInvariant().Contains(name) ||
                                    name.Contains(c.Name.Trim().ToLowerInvariant()));
                                if (fuzzy != null)
                                {
                                    categoryId = fuzzy.Id;
                                    categoryName = fuzzy.Name;
                                }
                            }
                        }
                    }
                }

                if (categoryId.HasValue)
                {
                    results.Add(new TransactionClassification(globalIndex, categoryId, categoryName, confidence));
                }
                else
                {
                    skippedNoMatch++;
                }
            }

            _logger.LogInformation(
                "Parsed {Matched} classifications from {Total} items. Skipped: {NoIndex} no index, {BadIndex} bad index, {LowConf} low confidence, {NoMatch} no category match",
                results.Count, classificationsArray.GetArrayLength(), skippedNoIndex, skippedBadIndex, skippedLowConf, skippedNoMatch);

            return results;
        }
        catch (JsonException ex)
        {
            _logger.LogWarning(ex, "Failed to parse Ollama JSON response: {Content}", responseContent.Length > 500 ? responseContent[..500] : responseContent);
            return new List<TransactionClassification>();
        }
    }

    public async Task<ClassifyUncategorizedResponse> ClassifyUncategorizedAsync(Guid userId, ClassifyUncategorizedRequest request)
    {
        var settings = await _db.AiCategorizationSettings.FirstOrDefaultAsync(s => s.UserId == userId);
        if (settings == null || !settings.IsEnabled)
            return new ClassifyUncategorizedResponse(0, 0, 0);

        var query = _db.Transactions
            .Include(t => t.Account)
            .Where(t => t.Account.UserId == userId && t.CategoryId == null);

        if (request.AccountId.HasValue)
            query = query.Where(t => t.AccountId == request.AccountId.Value);
        if (request.StartDate.HasValue)
            query = query.Where(t => t.Date >= request.StartDate.Value);
        if (request.EndDate.HasValue)
            query = query.Where(t => t.Date <= request.EndDate.Value);
        if (request.Status.HasValue)
            query = query.Where(t => t.Status == request.Status.Value);
        if (!string.IsNullOrEmpty(request.SearchText))
            query = query.Where(t =>
                (t.PayeeName != null && t.PayeeName.Contains(request.SearchText)) ||
                (t.Memo != null && t.Memo.Contains(request.SearchText)) ||
                (t.ReferenceNumber != null && t.ReferenceNumber.Contains(request.SearchText)));

        var uncategorized = await query.OrderByDescending(t => t.Date).ToListAsync();

        if (uncategorized.Count == 0)
            return new ClassifyUncategorizedResponse(0, 0, 0);

        var classifications = await ClassifyTransactionsAsync(userId, uncategorized);

        int classified = 0;
        foreach (var c in classifications)
        {
            if (c.CategoryId.HasValue && c.TransactionIndex >= 0 && c.TransactionIndex < uncategorized.Count)
            {
                uncategorized[c.TransactionIndex].CategoryId = c.CategoryId.Value;
                classified++;

                if (settings.UpdatePayeeDefaults)
                {
                    var txn = uncategorized[c.TransactionIndex];
                    if (txn.PayeeId.HasValue)
                    {
                        var payee = await _db.Payees.FindAsync(txn.PayeeId.Value);
                        if (payee != null && payee.DefaultCategoryId == null)
                            payee.DefaultCategoryId = c.CategoryId.Value;
                    }
                }
            }
        }

        await _db.SaveChangesAsync();

        return new ClassifyUncategorizedResponse(uncategorized.Count, classified, uncategorized.Count - classified);
    }
}