PlotDirector/PlotLine/Services/StoryIntelligenceJsonRepair.cs

using System.Text;
using System.Text.Json;
using System.Text.RegularExpressions;

namespace PlotLine.Services;

public static partial class StoryIntelligenceJsonRepair
{
    public static StoryIntelligenceJsonRepairResult Repair(string value)
    {
        var warnings = new List<string>();
        var extraction = ExtractJsonObject(value);
        if (!extraction.Success)
        {
            return StoryIntelligenceJsonRepairResult.Failed(value, extraction.ErrorMessage);
        }

        warnings.AddRange(extraction.Warnings);
        var numericRepaired = RepairTrailingDecimalNumbers(extraction.Json, warnings);
        var slashRepaired = RepairInvalidBackslashesInStrings(numericRepaired, warnings);

        if (!IsValidJson(slashRepaired, out var error))
        {
            return StoryIntelligenceJsonRepairResult.Failed(slashRepaired, error);
        }

        return new StoryIntelligenceJsonRepairResult(
            RepairedJson: slashRepaired,
            RepairApplied: warnings.Count > 0,
            Warnings: warnings,
            ErrorMessage: null);
    }

    public static StoryIntelligenceJsonExtractionResult ExtractJsonObject(string value)
    {
        var warnings = new List<string>();
        if (string.IsNullOrWhiteSpace(value))
        {
            return StoryIntelligenceJsonExtractionResult.Failed(value, "Output was empty.");
        }

        var candidate = StripMarkdownFence(value.Trim(), warnings);
        candidate = TrimToJsonObject(candidate, warnings);
        if (string.IsNullOrWhiteSpace(candidate))
        {
            return StoryIntelligenceJsonExtractionResult.Failed(value, "No complete JSON object was found.");
        }

        return new StoryIntelligenceJsonExtractionResult(
            Json: candidate,
            Warnings: warnings,
            ErrorMessage: null);
    }

    private static string StripMarkdownFence(string value, List<string> warnings)
    {
        var match = MarkdownFenceRegex().Match(value);
        if (!match.Success)
        {
            return value;
        }

        warnings.Add("Removed markdown code fence.");
        return match.Groups["json"].Value.Trim();
    }

    private static string TrimToJsonObject(string value, List<string> warnings)
    {
        var start = value.IndexOf('{');
        if (start < 0)
        {
            return string.Empty;
        }

        var depth = 0;
        var inString = false;
        var escaped = false;
        for (var i = start; i < value.Length; i++)
        {
            var ch = value[i];
            if (inString)
            {
                if (escaped)
                {
                    escaped = false;
                }
                else if (ch == '\\')
                {
                    escaped = true;
                }
                else if (ch == '"')
                {
                    inString = false;
                }

                continue;
            }

            if (ch == '"')
            {
                inString = true;
            }
            else if (ch == '{')
            {
                depth++;
            }
            else if (ch == '}')
            {
                depth--;
                if (depth == 0)
                {
                    if (start > 0 || i < value.Length - 1)
                    {
                        warnings.Add("Trimmed text outside the JSON object.");
                    }

                    return value[start..(i + 1)].Trim();
                }
            }
        }

        return string.Empty;
    }

    private static string RepairTrailingDecimalNumbers(string value, List<string> warnings)
    {
        var repaired = TrailingDecimalRegex().Replace(value, match =>
        {
            warnings.Add("Normalised trailing decimal number to valid JSON number.");
            return $"{match.Groups["prefix"].Value}{match.Groups["number"].Value}.0";
        });

        return repaired;
    }

    private static string RepairInvalidBackslashesInStrings(string value, List<string> warnings)
    {
        var builder = new StringBuilder(value.Length);
        var inString = false;
        var changed = false;

        for (var i = 0; i < value.Length; i++)
        {
            var ch = value[i];
            if (!inString)
            {
                builder.Append(ch);
                if (ch == '"')
                {
                    inString = true;
                }

                continue;
            }

            if (ch == '"')
            {
                inString = false;
                builder.Append(ch);
                continue;
            }

            if (ch == '\\')
            {
                var next = i + 1 < value.Length ? value[i + 1] : '\0';
                if (IsValidJsonEscape(next))
                {
                    builder.Append(ch);
                    if (next != '\0')
                    {
                        builder.Append(next);
                        i++;
                    }

                    continue;
                }

                builder.Append(@"\\");
                changed = true;
                continue;
            }

            builder.Append(ch);
        }

        if (changed)
        {
            warnings.Add("Escaped invalid backslash inside string value.");
        }

        return builder.ToString();
    }

    private static bool IsValidJsonEscape(char value)
        => value is '"' or '\\' or '/' or 'b' or 'f' or 'n' or 'r' or 't' or 'u';

    private static bool IsValidJson(string value, out string? error)
    {
        try
        {
            using var _ = JsonDocument.Parse(value);
            error = null;
            return true;
        }
        catch (JsonException ex)
        {
            error = ex.Message;
            return false;
        }
    }

    [GeneratedRegex("^```(?:json)?\\s*(?<json>.*?)\\s*```$", RegexOptions.Singleline | RegexOptions.IgnoreCase)]
    private static partial Regex MarkdownFenceRegex();

    [GeneratedRegex("(?<prefix>:\\s*)(?<number>-?\\d+)\\.\\s*(?=[,}\\]])")]
    private static partial Regex TrailingDecimalRegex();
}

public static class StoryIntelligenceAiJsonParser
{
    public static StoryIntelligenceAiJsonParseResult<T> Parse<T>(string rawAssistantOutput, JsonSerializerOptions jsonOptions)
    {
        var extraction = StoryIntelligenceJsonRepair.ExtractJsonObject(rawAssistantOutput);
        if (!extraction.Success)
        {
            return StoryIntelligenceAiJsonParseResult<T>.Failed(rawAssistantOutput, extraction.ErrorMessage ?? "No complete JSON object was found.");
        }

        try
        {
            var parsed = JsonSerializer.Deserialize<T>(extraction.Json, jsonOptions);
            return StoryIntelligenceAiJsonParseResult<T>.Ok(
                extraction.Json,
                repairedJson: null,
                parsed,
                extraction.Warnings);
        }
        catch (JsonException strictEx)
        {
            var repair = StoryIntelligenceJsonRepair.Repair(rawAssistantOutput);
            if (!repair.Success)
            {
                return StoryIntelligenceAiJsonParseResult<T>.Failed(extraction.Json, repair.ErrorMessage ?? strictEx.Message);
            }

            try
            {
                var parsed = JsonSerializer.Deserialize<T>(repair.RepairedJson, jsonOptions);
                var warnings = BuildRepairWarnings(repair.Warnings, strictEx.Path);
                return StoryIntelligenceAiJsonParseResult<T>.Ok(
                    extraction.Json,
                    repair.RepairedJson,
                    parsed,
                    warnings);
            }
            catch (JsonException repairedEx)
            {
                return StoryIntelligenceAiJsonParseResult<T>.Failed(extraction.Json, repairedEx.Message);
            }
        }
    }

    private static IReadOnlyList<string> BuildRepairWarnings(IReadOnlyList<string> repairWarnings, string? strictErrorPath)
    {
        var warnings = new List<string>();
        foreach (var warning in repairWarnings)
        {
            if (warning.Contains("trailing decimal", StringComparison.OrdinalIgnoreCase))
            {
                warnings.Add(string.IsNullOrWhiteSpace(strictErrorPath)
                    ? "Repaired malformed AI JSON numeric value."
                    : $"Repaired malformed AI JSON numeric value at {strictErrorPath}.");
            }
            else
            {
                warnings.Add(warning);
            }
        }

        return warnings;
    }
}

public sealed record StoryIntelligenceJsonRepairResult(
    string RepairedJson,
    bool RepairApplied,
    IReadOnlyList<string> Warnings,
    string? ErrorMessage)
{
    public bool Success => string.IsNullOrWhiteSpace(ErrorMessage);

    public static StoryIntelligenceJsonRepairResult Failed(string value, string? errorMessage)
        => new(value, false, [], errorMessage ?? "JSON repair failed.");
}

public sealed record StoryIntelligenceJsonExtractionResult(
    string Json,
    IReadOnlyList<string> Warnings,
    string? ErrorMessage)
{
    public bool Success => string.IsNullOrWhiteSpace(ErrorMessage);

    public static StoryIntelligenceJsonExtractionResult Failed(string value, string? errorMessage)
        => new(value, [], errorMessage ?? "JSON extraction failed.");
}

public sealed record StoryIntelligenceAiJsonParseResult<T>(
    bool Success,
    string RawJson,
    string? RepairedJson,
    T? Parsed,
    IReadOnlyList<string> Warnings,
    string? ErrorMessage)
{
    public static StoryIntelligenceAiJsonParseResult<T> Ok(
        string rawJson,
        string? repairedJson,
        T? parsed,
        IReadOnlyList<string> warnings)
        => new(true, rawJson, repairedJson, parsed, warnings, null);

    public static StoryIntelligenceAiJsonParseResult<T> Failed(string rawJson, string errorMessage)
        => new(false, rawJson, null, default, [], errorMessage);
}