h-school/src/HSchool.Server/Game/SwarmUiSettingsStore.cs

using System.Text.Json;
using System.Text.Json.Serialization;

namespace HSchool.Server.Game;

/// <summary>Loads, saves and serves SwarmUI generation presets from swarmui.json.</summary>
internal sealed class SwarmUiSettingsStore(IHostEnvironment environment, ILogger<SwarmUiSettingsStore> logger)
{
    private readonly Lock _lock = new();
    private SwarmUiConfigFile _config = SwarmUiConfigFile.CreateDefault();

    public string FilePath { get; } = Path.Combine(environment.ContentRootPath, "swarmui.json");

    public SwarmUiConfigFile Current
    {
        get
        {
            lock (_lock)
            {
                return _config;
            }
        }
    }

    public void Load()
    {
        lock (_lock)
        {
            _config = ReadFromDisk();
        }
    }

    public SwarmUiConfigFile Save(SwarmUiConfigFile config)
    {
        config.Validate();
        var json = JsonSerializer.Serialize(config, JsonOptions);
        var temp = FilePath + ".tmp";
        File.WriteAllText(temp, json);
        File.Move(temp, FilePath, overwrite: true);
        lock (_lock)
        {
            _config = config;
        }

        logger.LogInformation("SwarmUI settings saved to {Path}.", FilePath);
        return config;
    }

    public SwarmUiResolvedProfile Resolve(int age, PortraitKind kind)
    {
        lock (_lock)
        {
            return _config.Resolve(age, kind);
        }
    }

    private SwarmUiConfigFile ReadFromDisk()
    {
        if (!File.Exists(FilePath))
        {
            logger.LogWarning("SwarmUI settings file {Path} is missing; using built-in defaults.", FilePath);
            return SwarmUiConfigFile.CreateDefault();
        }

        try
        {
            var json = File.ReadAllText(FilePath);
            var loaded = JsonSerializer.Deserialize<SwarmUiConfigFile>(json, JsonOptions) ?? SwarmUiConfigFile.CreateDefault();
            loaded.NormalizeAfterLoad();
            loaded.Validate();
            return loaded;
        }
        catch (Exception ex)
        {
            logger.LogWarning(ex, "Could not read SwarmUI settings from {Path}; using built-in defaults.", FilePath);
            return SwarmUiConfigFile.CreateDefault();
        }
    }

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

internal sealed class SwarmUiConfigFile
{
    public string ActivePresetId { get; set; } = "default";

    public List<SwarmUiPresetDefinition> Presets { get; set; } = [];

    public List<SwarmUiAgeRule> AgeRules { get; set; } = [];

    // Legacy flat fields — read for migration only.
    public string? Model { get; set; }
    public int? Steps { get; set; }
    public double? CfgScale { get; set; }
    public int? ClipSkip { get; set; }
    public string? Sampler { get; set; }
    public string? Scheduler { get; set; }
    public long? Seed { get; set; }
    public string? Positive { get; set; }
    public string? Negative { get; set; }
    public SwarmUiKindPreset? Avatar { get; set; }
    public SwarmUiKindPreset? Custom { get; set; }
    public SwarmUiKindPreset? FullBody { get; set; }

    public void NormalizeAfterLoad()
    {
        if (Presets.Count == 0 && !string.IsNullOrWhiteSpace(Model))
        {
            Presets =
            [
                new SwarmUiPresetDefinition
                {
                    Id = "default",
                    Label = "Default",
                    Model = Model ?? "",
                    Steps = Steps ?? 8,
                    CfgScale = CfgScale ?? 1,
                    ClipSkip = ClipSkip ?? 1,
                    Sampler = Sampler ?? "",
                    Scheduler = Scheduler ?? "",
                    Seed = Seed ?? -1,
                    Positive = Positive ?? "",
                    Negative = Negative ?? "",
                    Avatar = Avatar ?? new SwarmUiKindPreset(),
                    Custom = Custom ?? new SwarmUiKindPreset(),
                    FullBody = FullBody ?? new SwarmUiKindPreset(),
                },
            ];
            ActivePresetId = "default";
        }

        if (Presets.Count == 0)
        {
            Presets = [SwarmUiPresetDefinition.CreateDefault()];
            ActivePresetId = "default";
        }

        foreach (var preset in Presets)
        {
            preset.Avatar ??= new SwarmUiKindPreset();
            preset.Custom ??= new SwarmUiKindPreset();
            preset.FullBody ??= new SwarmUiKindPreset();
            preset.PositiveLoras ??= [];
            preset.NegativeLoras ??= [];
        }

        if (string.IsNullOrWhiteSpace(ActivePresetId) || FindPreset(ActivePresetId) is null)
        {
            ActivePresetId = Presets[0].Id;
        }
    }

    public void Validate()
    {
        if (Presets.Count == 0)
        {
            throw new InvalidOperationException("At least one preset is required.");
        }

        var ids = new HashSet<string>(StringComparer.Ordinal);
        foreach (var preset in Presets)
        {
            if (string.IsNullOrWhiteSpace(preset.Id))
            {
                throw new InvalidOperationException("Every preset needs a non-empty id.");
            }

            if (!ids.Add(preset.Id))
            {
                throw new InvalidOperationException($"Duplicate preset id '{preset.Id}'.");
            }

            if (string.IsNullOrWhiteSpace(preset.Label))
            {
                preset.Label = preset.Id;
            }

            preset.Validate();
        }

        if (FindPreset(ActivePresetId) is null)
        {
            throw new InvalidOperationException($"Active preset '{ActivePresetId}' does not exist.");
        }

        foreach (var rule in AgeRules)
        {
            if (rule.MinAge > rule.MaxAge)
            {
                throw new InvalidOperationException($"Age rule {rule.MinAge}-{rule.MaxAge} is inverted.");
            }

            if (FindPreset(rule.PresetId) is null)
            {
                throw new InvalidOperationException($"Age rule references unknown preset '{rule.PresetId}'.");
            }
        }
    }

    public SwarmUiResolvedProfile Resolve(int age, PortraitKind kind)
    {
        var presetId = ResolvePresetId(age);
        var preset = FindPreset(presetId) ?? FindPreset(ActivePresetId) ?? Presets[0];
        return preset.ToProfile(kind);
    }

    private string ResolvePresetId(int age)
    {
        foreach (var rule in AgeRules.OrderBy(rule => rule.MinAge))
        {
            if (age >= rule.MinAge && age <= rule.MaxAge)
            {
                return rule.PresetId;
            }
        }

        return ActivePresetId;
    }

    private SwarmUiPresetDefinition? FindPreset(string id) =>
        Presets.FirstOrDefault(preset => string.Equals(preset.Id, id, StringComparison.Ordinal));

    public static SwarmUiConfigFile CreateDefault() =>
        new()
        {
            ActivePresetId = "default",
            Presets = [SwarmUiPresetDefinition.CreateDefault(), SwarmUiPresetDefinition.CreateChild()],
            AgeRules = [new SwarmUiAgeRule { MinAge = 6, MaxAge = 11, PresetId = "child" }],
        };

    public static SwarmUiConfigFile Clone(SwarmUiConfigFile source)
    {
        var json = JsonSerializer.Serialize(source, CloneJson);
        var copy = JsonSerializer.Deserialize<SwarmUiConfigFile>(json, CloneJson) ?? CreateDefault();
        copy.Model = null;
        copy.Steps = null;
        copy.CfgScale = null;
        copy.ClipSkip = null;
        copy.Sampler = null;
        copy.Scheduler = null;
        copy.Seed = null;
        copy.Positive = null;
        copy.Negative = null;
        copy.Avatar = null;
        copy.Custom = null;
        copy.FullBody = null;
        copy.NormalizeAfterLoad();
        return copy;
    }

    private static readonly JsonSerializerOptions CloneJson = new()
    {
        PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
        PropertyNameCaseInsensitive = true,
        DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,
    };
}

internal sealed class SwarmUiPresetDefinition
{
    public string Id { get; set; } = "";

    public string Label { get; set; } = "";

    public string Model { get; set; } = "";

    public int Steps { get; set; } = 8;

    public double CfgScale { get; set; } = 1;

    public int ClipSkip { get; set; } = 1;

    public string Sampler { get; set; } = "";

    public string Scheduler { get; set; } = "";

    public long Seed { get; set; } = -1;

    public string Positive { get; set; } = "";

    public string Negative { get; set; } = "";

    public List<SwarmUiLoraEntry>? PositiveLoras { get; set; }

    public List<SwarmUiLoraEntry>? NegativeLoras { get; set; }

    public SwarmUiKindPreset? Avatar { get; set; }

    public SwarmUiKindPreset? Custom { get; set; }

    public SwarmUiKindPreset? FullBody { get; set; }

    public void Validate()
    {
        if (Steps is < 1 or > 200)
        {
            throw new InvalidOperationException($"Preset '{Id}' steps must be between 1 and 200.");
        }

        if (CfgScale is < 0 or > 30)
        {
            throw new InvalidOperationException($"Preset '{Id}' cfgScale must be between 0 and 30.");
        }

        if (ClipSkip is < 0 or > 12)
        {
            throw new InvalidOperationException($"Preset '{Id}' clipSkip must be between 0 and 12.");
        }

        ValidateLoras(PositiveLoras, "positive");
        ValidateLoras(NegativeLoras, "negative");
    }

    private void ValidateLoras(IReadOnlyList<SwarmUiLoraEntry>? loras, string side)
    {
        if (loras is null)
        {
            return;
        }

        foreach (var lora in loras)
        {
            if (string.IsNullOrWhiteSpace(lora.Name))
            {
                throw new InvalidOperationException($"Preset '{Id}' has an empty {side} LoRA name.");
            }

            if (lora.Weight is < -4 or > 4)
            {
                throw new InvalidOperationException($"Preset '{Id}' LoRA '{lora.Name}' weight is out of range.");
            }
        }
    }

    public SwarmUiResolvedProfile ToProfile(PortraitKind kind) =>
        new(
            Id,
            Label,
            Model,
            Steps,
            CfgScale,
            ClipSkip,
            Sampler,
            Scheduler,
            Seed,
            Positive,
            Negative,
            PositiveLoras ?? [],
            NegativeLoras ?? [],
            KindPresetFor(kind));

    private SwarmUiKindPreset KindPresetFor(PortraitKind kind) => kind switch
    {
        PortraitKind.Avatar => Avatar ?? new SwarmUiKindPreset(),
        PortraitKind.Custom => Custom ?? new SwarmUiKindPreset(),
        PortraitKind.Full => FullBody ?? new SwarmUiKindPreset(),
        _ => throw new ArgumentOutOfRangeException(nameof(kind)),
    };

    public static SwarmUiPresetDefinition CreateDefault() =>
        new()
        {
            Id = "default",
            Label = "Default",
            Model = "DreamShaper_XL_-_Lightning_DPM++_SDE.safetensors",
            Steps = 4,
            CfgScale = 2,
            ClipSkip = 2,
            Sampler = "dpmpp_sde",
            Scheduler = "karras",
            Seed = 3346112079,
            Positive =
                "cinematic photo, realist detail, detailed character expressions, amazing quality, analog film grain, school portrait photograph, neutral background, natural lighting, realistic, sharp focus",
            Negative =
                "(low quality, worst quality:1.4), cgi, text, signature, watermark, extra limbs, nsfw, nude, naked, explicit, blurry, deformed, bad anatomy, logo",
            Avatar = new SwarmUiKindPreset
            {
                Width = 1024,
                Height = 1024,
                Positive = "close up, head and shoulders portrait, facing the camera, upper body visible.",
            },
            Custom = new SwarmUiKindPreset { Width = 896, Height = 1152 },
            FullBody = new SwarmUiKindPreset
            {
                Width = 896,
                Height = 1152,
                Positive = "full body standing portrait, head to toe visible, neutral pose, current outfit clearly visible.",
            },
        };

    public static SwarmUiPresetDefinition CreateChild()
    {
        var adult = CreateDefault();
        return new SwarmUiPresetDefinition
        {
            Id = "child",
            Label = "Children",
            Model = adult.Model,
            Steps = adult.Steps,
            CfgScale = adult.CfgScale,
            ClipSkip = adult.ClipSkip,
            Sampler = adult.Sampler,
            Scheduler = adult.Scheduler,
            Seed = adult.Seed,
            Positive =
                "cinematic photo, child-friendly school portrait, soft natural features, gentle expression, neutral background, natural lighting, realistic, sharp focus",
            Negative = adult.Negative,
            Avatar = new SwarmUiKindPreset
            {
                Width = 1024,
                Height = 1024,
                Positive = "close up, head and shoulders portrait of a child, facing the camera, upper body visible, soft features.",
            },
            Custom = adult.Custom,
            FullBody = new SwarmUiKindPreset
            {
                Width = 896,
                Height = 1152,
                Positive = "full body standing portrait of a child, head to toe visible, neutral pose, current outfit clearly visible.",
            },
        };
    }
}

internal sealed class SwarmUiAgeRule
{
    public int MinAge { get; set; }

    public int MaxAge { get; set; }

    public string PresetId { get; set; } = "";
}

internal sealed class SwarmUiKindPreset
{
    public int Width { get; set; } = 512;

    public int Height { get; set; } = 512;

    public string Positive { get; set; } = "";
}

internal sealed class SwarmUiLoraEntry
{
    public string Name { get; set; } = "";

    public double Weight { get; set; } = 1;
}

internal sealed record SwarmUiResolvedProfile(
    string PresetId,
    string PresetLabel,
    string Model,
    int Steps,
    double CfgScale,
    int ClipSkip,
    string Sampler,
    string Scheduler,
    long Seed,
    string Positive,
    string Negative,
    IReadOnlyList<SwarmUiLoraEntry> PositiveLoras,
    IReadOnlyList<SwarmUiLoraEntry> NegativeLoras,
    SwarmUiKindPreset KindPreset);