// Copyright (c) Microsoft. All rights reserved.

using System;
using System.Diagnostics.CodeAnalysis;
using Microsoft.Shared.DiagnosticIds;

namespace Microsoft.Agents.AI.Foundry.Hosting;

/// <summary>
/// Built-in <see cref="FoundryMemoryProvider"/> <c>stateInitializer</c> factories that derive the
/// <see cref="FoundryMemoryProviderScope"/> from the per-session <see cref="HostedSessionContext"/>
/// applied by the Foundry hosting layer.
/// </summary>
/// <remarks>
/// Pass the result of any of these helpers as the <c>stateInitializer</c> argument when constructing
/// <see cref="FoundryMemoryProvider"/>:
/// <code>
/// new FoundryMemoryProvider(client, "my-store",
///     stateInitializer: HostedFoundryMemoryProviderScopes.PerUser());
/// </code>
/// All helpers throw <see cref="InvalidOperationException"/> when
/// <see cref="HostedSessionContextExtensions.GetHostedContext"/> returns <see langword="null"/>.
/// That happens when the agent runs outside the Foundry hosting layer (e.g., a console app); in
/// that case write a custom <c>stateInitializer</c> instead of using these helpers.
/// </remarks>
[Experimental(DiagnosticIds.Experiments.AIOpenAIResponses)]
public static class HostedFoundryMemoryProviderScopes
{
    /// <summary>
    /// Returns a <c>stateInitializer</c> that scopes memories per end user, using
    /// <see cref="HostedSessionContext.UserId"/> as the partition key.
    /// </summary>
    /// <returns>A delegate suitable for the <c>stateInitializer</c> argument of <see cref="FoundryMemoryProvider"/>.</returns>
    public static Func<AgentSession?, FoundryMemoryProvider.State> PerUser() =>
        session => new FoundryMemoryProvider.State(new FoundryMemoryProviderScope(GetRequiredHostedContext(session).UserId));

    /// <summary>
    /// Returns a <c>stateInitializer</c> that scopes memories per conversation, using
    /// <see cref="HostedSessionContext.ChatId"/> as the partition key. Use this when memories should
    /// be visible to every participant in a shared conversation (for example, a Teams group chat).
    /// </summary>
    /// <returns>A delegate suitable for the <c>stateInitializer</c> argument of <see cref="FoundryMemoryProvider"/>.</returns>
    public static Func<AgentSession?, FoundryMemoryProvider.State> PerChat() =>
        session => new FoundryMemoryProvider.State(new FoundryMemoryProviderScope(GetRequiredHostedContext(session).ChatId));

    /// <summary>
    /// Returns a <c>stateInitializer</c> that scopes memories per (user, chat) pair, composing
    /// <see cref="HostedSessionContext.UserId"/> and <see cref="HostedSessionContext.ChatId"/> into a
    /// single delimiter-safe partition key. Use this when memories should be visible only to the same
    /// user within the same conversation.
    /// </summary>
    /// <remarks>
    /// Both identity values are opaque strings that may contain any characters, including the <c>:</c>
    /// delimiter. To keep the composite key injective (so two distinct (user, chat) pairs can never
    /// collide), each part is escaped (<c>\</c> becomes <c>\\</c>, then <c>:</c> becomes <c>\:</c>) before
    /// being joined with a <c>::</c> separator.
    /// </remarks>
    /// <returns>A delegate suitable for the <c>stateInitializer</c> argument of <see cref="FoundryMemoryProvider"/>.</returns>
    public static Func<AgentSession?, FoundryMemoryProvider.State> PerUserAndChat() =>
        session =>
        {
            var ctx = GetRequiredHostedContext(session);
            return new FoundryMemoryProvider.State(
                new FoundryMemoryProviderScope($"{EscapeScopePart(ctx.UserId)}::{EscapeScopePart(ctx.ChatId)}"));
        };

    /// <summary>
    /// Escapes special characters in a scope part so that distinct (user, chat) pairs produce distinct
    /// composite scope keys. Backslashes are escaped first (<c>\</c> becomes <c>\\</c>), then colons
    /// (<c>:</c> becomes <c>\:</c>), ensuring the <c>{user}::{chat}</c> format is unambiguous.
    /// </summary>
    private static string EscapeScopePart(string part) => part.Replace("\\", "\\\\").Replace(":", "\\:");

    private static HostedSessionContext GetRequiredHostedContext(AgentSession? session) =>
        session?.GetHostedContext()
            ?? throw new InvalidOperationException(
                $"{nameof(HostedSessionContext)} was not provided by the hosting layer. " +
                $"The {nameof(HostedFoundryMemoryProviderScopes)} helpers require the agent to be hosted via the Foundry hosting layer. " +
                "If running outside a hosted Foundry container, supply a custom stateInitializer to FoundryMemoryProvider instead.");
}