Wiki/src/Wiki.Application/Encyclopedia/AssignNotesUseCase.cs

using System.Net;
using System.Text;
using System.Text.Json.Nodes;
using Microsoft.Extensions.Logging;
using Wiki.Core.Model;
using Wiki.Core.Ports;
using static Wiki.Application.Encyclopedia.PlanParsing;

namespace Wiki.Application.Encyclopedia;

/// <summary>
/// <paramref name="New"/> are articles the batch introduces, with the batch notes that feed them as sources.
/// <paramref name="Assignments"/> maps batch notes to article slugs; an empty list is a skip.
/// <paramref name="Unusable"/>: the reply was truncated, not JSON, rejected, or held no valid entry.
/// </summary>
public sealed record AssignOutcome(
    IReadOnlyList<PlannedArticle> New,
    IReadOnlyDictionary<string, IReadOnlyList<string>> Assignments,
    bool Unusable);

/// <summary>Assigns one batch of notes to nearby existing articles or to new ones, in one bounded model call.</summary>
public sealed class AssignNotesUseCase(ITextModel model, EncyclopediaSettings settings, ILogger<AssignNotesUseCase> log)
{
    private const int FirstParagraphChars = 300;

    private static AssignOutcome Unusable() =>
        new([], new Dictionary<string, IReadOnlyList<string>>(), true);

    /// <summary>Model-unavailable and other errors propagate.</summary>
    public async Task<AssignOutcome> AssignAsync(
        IReadOnlyList<Document> batch, IReadOnlyList<PlannedArticle> candidates, ArticlePlan plan, CancellationToken ct)
    {
        var shown = candidates.Take(settings.CandidateArticles * 2).ToList();
        TextReply reply;
        try
        {
            reply = await model.CompleteAsync(Prompts.Assign, UserPrompt(batch, shown), settings.AssignMaxTokens, ct);
        }
        catch (HttpRequestException e) when (e.StatusCode == HttpStatusCode.BadRequest)
        {
            log.LogWarning(e, "Assign request was rejected");
            return Unusable();
        }

        if (reply.Truncated)
        {
            log.LogWarning("Assign reply hit the token cap ({Tokens})", settings.AssignMaxTokens);
            return Unusable();
        }

        var root = PlanParsing.Object(reply.Text);
        if (root is null)
        {
            log.LogWarning("Assign reply is not a JSON object");
            return Unusable();
        }

        var outcome = Validate(root, batch, shown, plan);
        if (outcome.Assignments.Count == 0)
        {
            log.LogWarning("Assign reply has no valid note entry");
            return Unusable();
        }
        return outcome;
    }

    private AssignOutcome Validate(JsonObject root, IReadOnlyList<Document> batch, List<PlannedArticle> candidates, ArticlePlan plan)
    {
        var candidateSlugs = candidates.Select(c => c.Slug).ToHashSet(StringComparer.Ordinal);
        var proposed = new Dictionary<string, PlannedArticle>(StringComparer.Ordinal);
        if (root["new"] is JsonArray items)
            foreach (var item in items)
            {
                if (item is not JsonObject o) continue;
                // a "new" slug that is a live article, or a redirect to one, is a reference to that article, not a new article
                if (Text(o["slug"]) is { } existing && plan.Resolve(existing) is { } live)
                {
                    log.LogInformation("Merged proposed new article {Slug} into the existing article {Live}", existing, live);
                    candidateSlugs.Add(live);
                    continue;
                }
                if (ReadNew(o, candidateSlugs, proposed.Keys, plan) is { } a) proposed[a.Slug] = a;
            }

        var inBatch = batch.Select(n => n.Slug).ToHashSet(StringComparer.Ordinal);
        var assignments = new Dictionary<string, IReadOnlyList<string>>(StringComparer.Ordinal);
        if (root["notes"] is JsonArray notes)
            foreach (var entry in notes)
            {
                if (entry is not JsonObject o || Text(o["note"]) is not { } note) continue;
                if (!inBatch.Contains(note) || assignments.ContainsKey(note))
                {
                    log.LogWarning("Dropping assignment entry for {Note}: not in the batch or duplicate", note);
                    continue;
                }
                if (o["skip"] is JsonValue s && s.TryGetValue<bool>(out var skip) && skip)
                {
                    assignments[note] = [];
                    continue;
                }
                var all = Strings(o["articles"]);
                var valid = all.Select(s => Referenced(s, candidateSlugs, plan))
                    .Where(s => candidateSlugs.Contains(s) || proposed.ContainsKey(s)).Distinct(StringComparer.Ordinal).ToList();
                if (valid.Count == 0)
                {
                    log.LogWarning("Dropping assignment entry for {Note}: no known article", note);
                    continue;
                }
                assignments[note] = valid;
            }

        var created = new List<PlannedArticle>();
        foreach (var a in proposed.Values)
        {
            var sources = assignments.Where(kv => kv.Value.Contains(a.Slug)).Select(kv => kv.Key).ToList();
            if (sources.Count == 0)
            {
                log.LogWarning("Dropping new article {Slug}: no note references it", a.Slug);
                continue;
            }
            created.Add(a with { Sources = sources });
        }
        return new AssignOutcome(created, assignments, false);
    }

    /// <summary>
    /// What a reference to <paramref name="slug"/> means: for a redirect to a live article that article, which becomes a valid
    /// reference; otherwise <paramref name="slug"/>.
    /// </summary>
    private static string Referenced(string slug, HashSet<string> candidateSlugs, ArticlePlan plan)
    {
        if (!plan.Redirects.ContainsKey(slug) || plan.Resolve(slug) is not { } live) return slug;
        candidateSlugs.Add(live);
        return live;
    }

    private PlannedArticle? ReadNew(JsonObject o, HashSet<string> candidateSlugs, IEnumerable<string> seen, ArticlePlan plan)
    {
        var slug = Text(o["slug"]);
        if (slug is null || !SlugPattern().IsMatch(slug))
        {
            log.LogWarning("Dropping new article with invalid slug {Slug}", slug);
            return null;
        }
        if (plan.SlugTaken(slug) || candidateSlugs.Contains(slug) || seen.Contains(slug))
        {
            log.LogWarning("Dropping new article {Slug}: slug taken or duplicate", slug);
            return null;
        }
        if (!TryCategory(Text(o["category"]), out var category))
        {
            log.LogWarning("Dropping new article {Slug}: unknown category", slug);
            return null;
        }
        var scope = Text(o["scope"]);
        if (scope is null)
        {
            log.LogWarning("Dropping new article {Slug}: empty scope", slug);
            return null;
        }
        if (Text(o["titleDe"]) is not { } titleDe || Text(o["titleEn"]) is not { } titleEn)
        {
            log.LogWarning("Dropping new article {Slug}: missing title", slug);
            return null;
        }
        return new PlannedArticle(slug, titleDe, titleEn, category, scope, []);
    }

    private static string UserPrompt(IReadOnlyList<Document> batch, List<PlannedArticle> candidates)
    {
        var sb = new StringBuilder();
        sb.Append("Notes (slug | title | tags | body length in chars | first paragraph):\n");
        foreach (var n in batch.OrderBy(n => n.Slug, StringComparer.Ordinal))
            sb.Append(n.Slug).Append(" | ").Append(Flatten(n.Title)).Append(" | ").Append(Flatten(string.Join(", ", n.Tags)))
              .Append(" | ").Append(n.Body.Length).Append(" | ").Append(FirstParagraph(n.Body)).Append('\n');
        sb.Append("\nCandidate articles (slug | title | category | scope | source count):\n");
        foreach (var c in candidates)
            sb.Append(c.Slug).Append(" | ").Append(Flatten(c.TitleDe)).Append(" | ").Append(c.Category.ToString().ToLowerInvariant())
              .Append(" | ").Append(Flatten(c.Scope)).Append(" | ").Append(c.Sources.Count).Append('\n');
        return sb.ToString();
    }

    /// <summary>First non-empty paragraph after leading heading lines, flattened and cut.</summary>
    private static string FirstParagraph(string body)
    {
        var lines = body.Replace("\r", "").Split('\n')
            .SkipWhile(l => string.IsNullOrWhiteSpace(l) || l.TrimStart().StartsWith('#'))
            .TakeWhile(l => !string.IsNullOrWhiteSpace(l));
        var text = Flatten(string.Join(' ', lines)).Trim();
        return text.Length > FirstParagraphChars ? text[..FirstParagraphChars] : text;
    }
}