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

using HSchool.Server.Api;

namespace HSchool.Server.Game;

internal sealed class PortraitService(
    SchoolStore store,
    SwarmUiClient swarm,
    SwarmUiSettingsStore settingsStore,
    GameLoopService loop,
    GameCommandQueue commands,
    ILogger<PortraitService> logger)
{
    public const int MaxCustomPromptLength = 2000;

    private static readonly TimeSpan PersonLookupTimeout = TimeSpan.FromSeconds(5);

    /// <summary>English labels for Swarm prompts, independent of the UI locale.</summary>
    private const string PromptLocale = "en";

    public bool IsGenerationEnabled => swarm.IsConfigured;

    public (bool HasAvatar, bool HasCustom, bool HasFullBody) Flags(int schoolId, string personId) =>
        store.PortraitFlags(schoolId, personId);

    public PersonCardResponse WithPortraitFlags(int schoolId, PersonCardResponse card)
    {
        var (hasAvatar, hasCustom, hasFullBody) = Flags(schoolId, card.Id);
        var customPrompt = hasCustom ? store.TryReadCustomPortraitPrompt(schoolId, card.Id) : null;
        return card with
        {
            HasAvatar = hasAvatar,
            HasCustom = hasCustom,
            HasFullBody = hasFullBody,
            CustomPortraitPrompt = customPrompt,
        };
    }

    public async Task<PersonLookupError> EnsurePersonAsync(
        int schoolId,
        string personId,
        CancellationToken cancellationToken)
    {
        var command = new GameCommand.GetPerson(
            schoolId,
            personId,
            PromptLocale,
            NewCompletion<PersonCardResult>());
        commands.Enqueue(command);

        var outcome = await command.Result.Task.WaitAsync(PersonLookupTimeout, cancellationToken);
        return outcome.Error;
    }

    public async Task<PortraitPromptBuildResult> BuildPromptAsync(
        int schoolId,
        string personId,
        PortraitKind kind,
        string? promptExtra,
        CancellationToken cancellationToken)
    {
        var resolved = ResolveCustomPromptExtra(schoolId, personId, kind, promptExtra);
        if (kind == PortraitKind.Custom && resolved is null)
        {
            return PortraitPromptBuildResult.InvalidPrompt;
        }

        var outcome = await LookupBuildInputAsync(schoolId, personId, cancellationToken);
        if (outcome.Error == PersonLookupError.UnknownPerson)
        {
            return PortraitPromptBuildResult.UnknownPerson;
        }

        if (outcome.Error != PersonLookupError.None || outcome.Card is null)
        {
            return PortraitPromptBuildResult.UnknownSchool;
        }

        var profile = ConfigFor(schoolId).Resolve(outcome.Card.Age, kind);
        var (positive, negative) = PortraitPromptBuilder.Build(
            outcome.Card,
            profile,
            kind,
            resolved,
            outcome.Scene);
        return PortraitPromptBuildResult.Succeeded(
            kind,
            positive,
            negative,
            kind == PortraitKind.Custom ? resolved : null,
            profile.PresetId,
            profile.PresetLabel);
    }

    public async Task<PortraitGenerationResult> GenerateAsync(
        int schoolId,
        string personId,
        PortraitKind kind,
        string? promptExtra,
        CancellationToken cancellationToken)
    {
        if (!swarm.IsConfigured)
        {
            return PortraitGenerationResult.NotConfigured;
        }

        if (kind == PortraitKind.Custom)
        {
            promptExtra = promptExtra?.Trim();
            if (string.IsNullOrWhiteSpace(promptExtra))
            {
                return PortraitGenerationResult.InvalidPrompt;
            }

            if (promptExtra.Length > MaxCustomPromptLength)
            {
                return PortraitGenerationResult.InvalidPrompt;
            }
        }

        var outcome = await LookupBuildInputAsync(schoolId, personId, cancellationToken);
        if (outcome.Error == PersonLookupError.UnknownPerson)
        {
            return PortraitGenerationResult.UnknownPerson;
        }

        if (outcome.Error != PersonLookupError.None || outcome.Card is null)
        {
            return PortraitGenerationResult.UnknownSchool;
        }

        var profile = ConfigFor(schoolId).Resolve(outcome.Card.Age, kind);
        var (positive, negative) = PortraitPromptBuilder.Build(
            outcome.Card,
            profile,
            kind,
            promptExtra,
            outcome.Scene);

        try
        {
            var bytes = await swarm.GenerateAsync(
                positive,
                negative,
                PortraitPromptBuilder.ApplyScene(profile, outcome.Scene),
                cancellationToken);
            store.SavePortrait(schoolId, personId, kind, bytes, kind == PortraitKind.Custom ? promptExtra : null);
            var flags = Flags(schoolId, personId);
            var savedPrompt = kind == PortraitKind.Custom ? promptExtra : store.TryReadCustomPortraitPrompt(schoolId, personId);
            return PortraitGenerationResult.Succeeded(
                kind,
                flags.HasAvatar,
                flags.HasCustom,
                flags.HasFullBody,
                savedPrompt);
        }
        catch (TaskCanceledException ex) when (!cancellationToken.IsCancellationRequested)
        {
            logger.LogWarning(ex, "SwarmUI timed out for school {SchoolId} person {PersonId}.", schoolId, personId);
            return PortraitGenerationResult.TimedOut;
        }
        catch (HttpRequestException ex)
        {
            logger.LogWarning(ex, "SwarmUI request failed for school {SchoolId} person {PersonId}.", schoolId, personId);
            return PortraitGenerationResult.Unavailable;
        }
        catch (Exception ex)
        {
            logger.LogWarning(ex, "Portrait generation failed for school {SchoolId} person {PersonId}.", schoolId, personId);
            return PortraitGenerationResult.Unavailable;
        }
    }

    private SwarmUiConfigFile ConfigFor(int schoolId) =>
        loop.PortraitSettingsOf(schoolId) ?? settingsStore.Current;

    private async Task<PortraitBuildInputResult> LookupBuildInputAsync(
        int schoolId,
        string personId,
        CancellationToken cancellationToken)
    {
        var command = new GameCommand.GetPortraitBuildInput(
            schoolId,
            personId,
            PromptLocale,
            NewCompletion<PortraitBuildInputResult>());
        commands.Enqueue(command);
        return await command.Result.Task.WaitAsync(PersonLookupTimeout, cancellationToken);
    }

    private async Task<PersonCardResult> LookupPersonAsync(
        int schoolId,
        string personId,
        CancellationToken cancellationToken)
    {
        var command = new GameCommand.GetPerson(
            schoolId,
            personId,
            PromptLocale,
            NewCompletion<PersonCardResult>());
        commands.Enqueue(command);
        return await command.Result.Task.WaitAsync(PersonLookupTimeout, cancellationToken);
    }

    private string? ResolveCustomPromptExtra(int schoolId, string personId, PortraitKind kind, string? promptExtra)
    {
        if (kind != PortraitKind.Custom)
        {
            return null;
        }

        promptExtra = promptExtra?.Trim();
        if (string.IsNullOrWhiteSpace(promptExtra))
        {
            promptExtra = store.TryReadCustomPortraitPrompt(schoolId, personId)?.Trim();
        }

        if (string.IsNullOrWhiteSpace(promptExtra) || promptExtra.Length > MaxCustomPromptLength)
        {
            return null;
        }

        return promptExtra;
    }

    private static TaskCompletionSource<T> NewCompletion<T>() =>
        new(TaskCreationOptions.RunContinuationsAsynchronously);
}

internal enum PortraitPromptBuildOutcome
{
    Succeeded,
    UnknownSchool,
    UnknownPerson,
    InvalidPrompt,
}

internal sealed record PortraitPromptBuildResult(
    PortraitPromptBuildOutcome Outcome,
    PortraitKind Kind,
    string Positive,
    string Negative,
    string? PromptExtra,
    string PresetId,
    string PresetLabel)
{
    public static PortraitPromptBuildResult UnknownSchool { get; } =
        new(PortraitPromptBuildOutcome.UnknownSchool, default, string.Empty, string.Empty, null, "", "");

    public static PortraitPromptBuildResult UnknownPerson { get; } =
        new(PortraitPromptBuildOutcome.UnknownPerson, default, string.Empty, string.Empty, null, "", "");

    public static PortraitPromptBuildResult InvalidPrompt { get; } =
        new(PortraitPromptBuildOutcome.InvalidPrompt, default, string.Empty, string.Empty, null, "", "");

    public static PortraitPromptBuildResult Succeeded(
        PortraitKind kind,
        string positive,
        string negative,
        string? promptExtra,
        string presetId,
        string presetLabel) =>
        new(PortraitPromptBuildOutcome.Succeeded, kind, positive, negative, promptExtra, presetId, presetLabel);
}

internal sealed record PortraitPromptResponse(
    string Kind,
    string Positive,
    string Negative,
    string? PromptExtra,
    string PresetId,
    string PresetLabel);

internal enum PortraitGenerationOutcome
{
    Succeeded,
    UnknownSchool,
    UnknownPerson,
    NotConfigured,
    InvalidPrompt,
    Unavailable,
    TimedOut,
}

internal sealed record PortraitGenerationResult(
    PortraitGenerationOutcome Outcome,
    PortraitKind Kind,
    bool HasAvatar,
    bool HasCustom,
    bool HasFullBody,
    string? CustomPortraitPrompt)
{
    public static PortraitGenerationResult UnknownSchool { get; } =
        new(PortraitGenerationOutcome.UnknownSchool, default, false, false, false, null);

    public static PortraitGenerationResult UnknownPerson { get; } =
        new(PortraitGenerationOutcome.UnknownPerson, default, false, false, false, null);

    public static PortraitGenerationResult NotConfigured { get; } =
        new(PortraitGenerationOutcome.NotConfigured, default, false, false, false, null);

    public static PortraitGenerationResult InvalidPrompt { get; } =
        new(PortraitGenerationOutcome.InvalidPrompt, default, false, false, false, null);

    public static PortraitGenerationResult Unavailable { get; } =
        new(PortraitGenerationOutcome.Unavailable, default, false, false, false, null);

    public static PortraitGenerationResult TimedOut { get; } =
        new(PortraitGenerationOutcome.TimedOut, default, false, false, false, null);

    public static PortraitGenerationResult Succeeded(
        PortraitKind kind,
        bool hasAvatar,
        bool hasCustom,
        bool hasFullBody,
        string? customPortraitPrompt) =>
        new(PortraitGenerationOutcome.Succeeded, kind, hasAvatar, hasCustom, hasFullBody, customPortraitPrompt);
}

internal sealed record GeneratePortraitRequest(string? PromptExtra);

internal sealed record PortraitResponse(
    string Kind,
    bool HasAvatar,
    bool HasCustom,
    bool HasFullBody,
    string? CustomPortraitPrompt);