PkmnLib.NET/Plugins/PkmnLib.Plugin.Gen7.SourceGen/Parsing/DataExtractor.cs

using System.Text.Json;

namespace PkmnLib.Plugin.Gen7.SourceGen.Parsing;

/// <summary>
/// A single generated constants class: its name and the raw data names it should contain.
/// </summary>
public sealed record ClassSpec(string ClassName, string SourceFile, EquatableArray<string> Names);

/// <summary>
/// A diagnostic produced during extraction, kept as plain data so pipeline values stay equatable.
/// </summary>
public sealed record ExtractionDiagnostic(string Id, string Message);

/// <summary>
/// The result of extracting all constant classes from a single data file.
/// </summary>
public sealed record ExtractionResult(
    EquatableArray<ClassSpec> Classes,
    EquatableArray<ExtractionDiagnostic> Diagnostics);

/// <summary>
/// Extracts the names for the generated constant classes from the Gen 7 data files.
/// </summary>
public static class DataExtractor
{
    private static readonly HashSet<string> KnownFiles = new(StringComparer.Ordinal)
    {
        "Moves.jsonc", "Abilities.jsonc", "Items.json", "Pokemon.json", "Types.csv", "Natures.csv",
    };

    /// <summary>
    /// Matches the runtime's <c>JsonOptions</c>: the data files contain <c>//</c> comments
    /// (including Items.json, despite its extension) and may contain trailing commas.
    /// </summary>
    private static readonly JsonDocumentOptions JsonOptions = new()
    {
        CommentHandling = JsonCommentHandling.Skip,
        AllowTrailingCommas = true,
    };

    public static bool IsKnownFile(string fileName) => KnownFiles.Contains(fileName);

    /// <summary>
    /// Extracts the constant classes for a data file. Never throws: a generator that throws is
    /// silently disabled by IDE hosts, so failures are reported as diagnostics instead.
    /// </summary>
    public static ExtractionResult Extract(string fileName, string? text)
    {
        if (text is null)
        {
            return Error(fileName, "file content could not be read");
        }
        try
        {
            return fileName switch
            {
                "Moves.jsonc" => ExtractMoves(fileName, text),
                "Abilities.jsonc" => ExtractTopLevelKeys(fileName, text, "AbilityNames"),
                "Items.json" => ExtractItems(fileName, text),
                "Pokemon.json" => ExtractSpecies(fileName, text),
                "Types.csv" => ExtractTypes(fileName, text),
                "Natures.csv" => ExtractNatures(fileName, text),
                _ => new ExtractionResult(new EquatableArray<ClassSpec>([]),
                    new EquatableArray<ExtractionDiagnostic>([])),
            };
        }
        catch (Exception e)
        {
            return Error(fileName, e.Message);
        }
    }

    private static ExtractionResult Error(string fileName, string message) =>
        new(new EquatableArray<ClassSpec>([]), new EquatableArray<ExtractionDiagnostic>([
            new ExtractionDiagnostic("PKMN7G000", $"Failed to parse data file '{fileName}': {message}"),
        ]));

    private static ExtractionResult ExtractMoves(string fileName, string text)
    {
        var moveNames = new HashSet<string>(StringComparer.Ordinal);
        var flags = new HashSet<string>(StringComparer.Ordinal);
        var effectNames = new HashSet<string>(StringComparer.Ordinal);
        using var doc = JsonDocument.Parse(text, JsonOptions);
        foreach (var move in doc.RootElement.GetProperty("data").EnumerateArray())
        {
            Add(moveNames, move.GetProperty("name").GetString());
            if (move.TryGetProperty("flags", out var flagsElement))
            {
                foreach (var flag in flagsElement.EnumerateArray())
                {
                    Add(flags, flag.GetString());
                }
            }
            if (move.TryGetProperty("effect", out var effect) && effect.TryGetProperty("name", out var effectName))
            {
                Add(effectNames, effectName.GetString());
            }
        }
        return Result(Class("MoveNames", fileName, moveNames), Class("MoveFlags", fileName, flags),
            Class("MoveEffectNames", fileName, effectNames));
    }

    private static ExtractionResult ExtractTopLevelKeys(string fileName, string text, string className)
    {
        var names = new HashSet<string>(StringComparer.Ordinal);
        using var doc = JsonDocument.Parse(text, JsonOptions);
        foreach (var property in doc.RootElement.EnumerateObject())
        {
            if (!property.Name.StartsWith("$", StringComparison.Ordinal))
                Add(names, property.Name);
        }
        return Result(Class(className, fileName, names));
    }

    private static ExtractionResult ExtractItems(string fileName, string text)
    {
        var names = new HashSet<string>(StringComparer.Ordinal);
        using var doc = JsonDocument.Parse(text, JsonOptions);
        foreach (var item in doc.RootElement.EnumerateArray())
        {
            Add(names, item.GetProperty("name").GetString());
        }
        return Result(Class("ItemNames", fileName, names));
    }

    private static ExtractionResult ExtractSpecies(string fileName, string text)
    {
        var speciesNames = new HashSet<string>(StringComparer.Ordinal);
        var formNames = new HashSet<string>(StringComparer.Ordinal);
        using var doc = JsonDocument.Parse(text, JsonOptions);
        foreach (var property in doc.RootElement.EnumerateObject())
        {
            if (property.Name.StartsWith("$", StringComparison.Ordinal))
                continue;
            Add(speciesNames, property.Name);
            if (property.Value.TryGetProperty("formes", out var formes))
            {
                foreach (var forme in formes.EnumerateObject())
                {
                    Add(formNames, forme.Name);
                }
            }
        }
        return Result(Class("SpeciesNames", fileName, speciesNames), Class("FormNames", fileName, formNames));
    }

    private static ExtractionResult ExtractTypes(string fileName, string text)
    {
        var names = new HashSet<string>(StringComparer.Ordinal);
        var header = FirstLine(text);
        var cells = header.Split(header.Contains('|') ? '|' : ',');
        for (var i = 1; i < cells.Length; i++)
        {
            Add(names, cells[i]);
        }
        return Result(Class("TypeNames", fileName, names));
    }

    private static ExtractionResult ExtractNatures(string fileName, string text)
    {
        var names = new HashSet<string>(StringComparer.Ordinal);
        var lines = text.Split('\n');
        for (var i = 1; i < lines.Length; i++)
        {
            var line = lines[i].Trim();
            if (line.Length == 0)
                continue;
            Add(names, line.Split(line.Contains('|') ? '|' : ',')[0]);
        }
        return Result(Class("NatureNames", fileName, names));
    }

    private static string FirstLine(string text)
    {
        var index = text.IndexOf('\n');
        return (index < 0 ? text : text.Substring(0, index)).Trim();
    }

    private static void Add(HashSet<string> names, string? raw)
    {
        var name = raw?.Trim().ToLowerInvariant();
        if (!string.IsNullOrEmpty(name))
            names.Add(name!);
    }

    private static ClassSpec Class(string className, string fileName, HashSet<string> names)
    {
        var sorted = names.ToArray();
        Array.Sort(sorted, StringComparer.Ordinal);
        return new ClassSpec(className, fileName, new EquatableArray<string>(sorted));
    }

    private static ExtractionResult Result(params ClassSpec[] classes) =>
        new(new EquatableArray<ClassSpec>(classes), new EquatableArray<ExtractionDiagnostic>([]));
}