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

using System;
using System.Diagnostics.CodeAnalysis;
using A2A;
using Microsoft.Agents.AI;
using Microsoft.Agents.AI.Hosting;
using Microsoft.Agents.AI.Hosting.A2A;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Shared.DiagnosticIds;

namespace Microsoft.Extensions.DependencyInjection;

/// <summary>
/// Provides extension methods for registering A2A server instances in the dependency injection container.
/// </summary>
[Experimental(DiagnosticIds.Experiments.AIResponseContinuations)]
public static class A2AServerServiceCollectionExtensions
{
    /// <summary>
    /// Registers an <see cref="A2AServer"/> in the dependency injection container, keyed by the agent name
    /// specified in the <paramref name="agentBuilder"/>. This method only registers the server; to expose it
    /// as an HTTP endpoint, call one of the <c>MapA2AHttpJson</c> or <c>MapA2AJsonRpc</c> endpoint mapping
    /// methods during application startup.
    /// </summary>
    /// <param name="agentBuilder">The agent builder whose name identifies the agent.</param>
    /// <param name="configureOptions">An optional callback to configure <see cref="A2AServerRegistrationOptions"/>.</param>
    /// <returns>The <paramref name="agentBuilder"/> for chaining.</returns>
    /// <remarks>
    /// <para>
    /// <strong>Trust model.</strong> The A2A <c>contextId</c> arrives from the wire
    /// and is treated as a chain-resume identifier — <em>not</em> as an authorization
    /// token. The <see cref="AgentSessionStore"/> contract carries no principal/owner
    /// dimension, so when a persistent store is registered any caller who knows or
    /// guesses another caller's <c>contextId</c> can resume that other caller's
    /// persisted thread. Hosts that serve more than one user must compose a principal
    /// dimension into the lookup key — typically by calling
    /// <c>UseClaimsBasedSessionIsolation(...)</c> from
    /// <c>Microsoft.Agents.AI.Hosting.AspNetCore</c> (or by registering a custom
    /// <see cref="SessionIsolationKeyProvider"/>). When no isolation provider is
    /// registered, behavior is unchanged — the bare <c>contextId</c> is used as the
    /// conversation identifier, which is appropriate for first-run / single-user /
    /// prototyping scenarios but unsafe for multi-user hosts.
    /// </para>
    /// </remarks>
    public static IHostedAgentBuilder AddA2AServer(this IHostedAgentBuilder agentBuilder, Action<A2AServerRegistrationOptions>? configureOptions = null)
    {
        ArgumentNullException.ThrowIfNull(agentBuilder);

        agentBuilder.ServiceCollection.AddA2AServer(agentBuilder.Name, configureOptions);

        return agentBuilder;
    }

    /// <summary>
    /// Registers an <see cref="A2AServer"/> in the dependency injection container, keyed by the specified
    /// agent name. This method only registers the server; to expose it as an HTTP endpoint, call one of the
    /// <c>MapA2AHttpJson</c> or <c>MapA2AJsonRpc</c> endpoint mapping methods during application startup.
    /// </summary>
    /// <param name="builder">The host application builder to configure.</param>
    /// <param name="agentName">The name of the agent to create an A2A server for.</param>
    /// <param name="configureOptions">An optional callback to configure <see cref="A2AServerRegistrationOptions"/>.</param>
    /// <returns>The <paramref name="builder"/> for chaining.</returns>
    /// <remarks>
    /// See the trust-model remarks on <see cref="AddA2AServer(IHostedAgentBuilder, Action{A2AServerRegistrationOptions}?)"/>
    /// for guidance on multi-user hosts (the wire <c>contextId</c> is a chain-resume
    /// identifier, not an authorization token; multi-user hosts must compose a
    /// principal dimension via <c>UseClaimsBasedSessionIsolation(...)</c> or a custom
    /// <see cref="SessionIsolationKeyProvider"/>).
    /// </remarks>
    public static IHostApplicationBuilder AddA2AServer(this IHostApplicationBuilder builder, string agentName, Action<A2AServerRegistrationOptions>? configureOptions = null)
    {
        ArgumentNullException.ThrowIfNull(builder);

        builder.Services.AddA2AServer(agentName, configureOptions);

        return builder;
    }

    /// <summary>
    /// Registers an <see cref="A2AServer"/> in the dependency injection container for the specified
    /// <see cref="AIAgent"/> instance, keyed by the agent's <see cref="AIAgent.Name"/>. This method only
    /// registers the server; to expose it as an HTTP endpoint, call one of the <c>MapA2AHttpJson</c> or
    /// <c>MapA2AJsonRpc</c> endpoint mapping methods during application startup.
    /// </summary>
    /// <param name="builder">The host application builder to configure.</param>
    /// <param name="agent">The agent instance to create an A2A server for.</param>
    /// <param name="configureOptions">An optional callback to configure <see cref="A2AServerRegistrationOptions"/>.</param>
    /// <returns>The <paramref name="builder"/> for chaining.</returns>
    /// <remarks>
    /// See the trust-model remarks on <see cref="AddA2AServer(IHostedAgentBuilder, Action{A2AServerRegistrationOptions}?)"/>
    /// for guidance on multi-user hosts (the wire <c>contextId</c> is a chain-resume
    /// identifier, not an authorization token; multi-user hosts must compose a
    /// principal dimension via <c>UseClaimsBasedSessionIsolation(...)</c> or a custom
    /// <see cref="SessionIsolationKeyProvider"/>).
    /// </remarks>
    public static IHostApplicationBuilder AddA2AServer(this IHostApplicationBuilder builder, AIAgent agent, Action<A2AServerRegistrationOptions>? configureOptions = null)
    {
        ArgumentNullException.ThrowIfNull(builder);

        builder.Services.AddA2AServer(agent, configureOptions);

        return builder;
    }

    /// <summary>
    /// Registers an <see cref="A2AServer"/> in the dependency injection container, keyed by the specified
    /// agent name. This method only registers the server; to expose it as an HTTP endpoint, call one of the
    /// <c>MapA2AHttpJson</c> or <c>MapA2AJsonRpc</c> endpoint mapping methods during application startup.
    /// </summary>
    /// <param name="services">The service collection to add the A2A server to.</param>
    /// <param name="agentName">The name of the agent to create an A2A server for.</param>
    /// <param name="configureOptions">An optional callback to configure <see cref="A2AServerRegistrationOptions"/>.</param>
    /// <returns>The <paramref name="services"/> for chaining.</returns>
    /// <remarks>
    /// See the trust-model remarks on <see cref="AddA2AServer(IHostedAgentBuilder, Action{A2AServerRegistrationOptions}?)"/>
    /// for guidance on multi-user hosts (the wire <c>contextId</c> is a chain-resume
    /// identifier, not an authorization token; multi-user hosts must compose a
    /// principal dimension via <c>UseClaimsBasedSessionIsolation(...)</c> or a custom
    /// <see cref="SessionIsolationKeyProvider"/>).
    /// </remarks>
    public static IServiceCollection AddA2AServer(this IServiceCollection services, string agentName, Action<A2AServerRegistrationOptions>? configureOptions = null)
    {
        ArgumentNullException.ThrowIfNull(services);
        ArgumentException.ThrowIfNullOrWhiteSpace(agentName);

        A2AServerRegistrationOptions? options = null;
        if (configureOptions is not null)
        {
            options = new A2AServerRegistrationOptions();
            configureOptions(options);
        }

        services.AddKeyedSingleton(agentName, (sp, _) =>
        {
            var agent = sp.GetRequiredKeyedService<AIAgent>(agentName);
            return CreateA2AServer(sp, agent, options);
        });

        return services;
    }

    /// <summary>
    /// Registers an <see cref="A2AServer"/> in the dependency injection container for the specified
    /// <see cref="AIAgent"/> instance, keyed by the agent's <see cref="AIAgent.Name"/>. This method only
    /// registers the server; to expose it as an HTTP endpoint, call one of the <c>MapA2AHttpJson</c> or
    /// <c>MapA2AJsonRpc</c> endpoint mapping methods during application startup.
    /// </summary>
    /// <param name="services">The service collection to add the A2A server to.</param>
    /// <param name="agent">The agent instance to create an A2A server for.</param>
    /// <param name="configureOptions">An optional callback to configure <see cref="A2AServerRegistrationOptions"/>.</param>
    /// <returns>The <paramref name="services"/> for chaining.</returns>
    /// <remarks>
    /// See the trust-model remarks on <see cref="AddA2AServer(IHostedAgentBuilder, Action{A2AServerRegistrationOptions}?)"/>
    /// for guidance on multi-user hosts (the wire <c>contextId</c> is a chain-resume
    /// identifier, not an authorization token; multi-user hosts must compose a
    /// principal dimension via <c>UseClaimsBasedSessionIsolation(...)</c> or a custom
    /// <see cref="SessionIsolationKeyProvider"/>).
    /// </remarks>
    public static IServiceCollection AddA2AServer(this IServiceCollection services, AIAgent agent, Action<A2AServerRegistrationOptions>? configureOptions = null)
    {
        ArgumentNullException.ThrowIfNull(services);
        ArgumentNullException.ThrowIfNull(agent);
        ArgumentException.ThrowIfNullOrWhiteSpace(agent.Name, nameof(agent) + "." + nameof(agent.Name));

        A2AServerRegistrationOptions? options = null;
        if (configureOptions is not null)
        {
            options = new A2AServerRegistrationOptions();
            configureOptions(options);
        }

        services.AddKeyedSingleton(agent.Name, (sp, _) => CreateA2AServer(sp, agent, options));

        return services;
    }

    private static A2AServer CreateA2AServer(IServiceProvider serviceProvider, AIAgent agent, A2AServerRegistrationOptions? options)
    {
        var agentHandler = serviceProvider.GetKeyedService<IAgentHandler>(agent.Name);
        if (agentHandler is null)
        {
            var agentSessionStore = serviceProvider.GetKeyedService<AgentSessionStore>(agent.Name);
            var runMode = options?.AgentRunMode ?? AgentRunMode.DisallowBackground;

            // Ensure that we have an IsolationKeyScopedAgentSessionStore registered.
            var isolationKeyProvider = serviceProvider.GetService<SessionIsolationKeyProvider>();
            if (agentSessionStore?.GetService<IsolationKeyScopedAgentSessionStore>() is null)
            {
                agentSessionStore ??= new InMemoryAgentSessionStore();
                agentSessionStore = new IsolationKeyScopedAgentSessionStore(agentSessionStore, isolationKeyProvider, new() { Strict = isolationKeyProvider != null });
            }

            var hostAgent = new AIHostAgent(
                innerAgent: agent,
                sessionStore: agentSessionStore);

            agentHandler = new A2AAgentHandler(hostAgent, runMode);
        }

        var loggerFactory = serviceProvider.GetService<ILoggerFactory>() ?? NullLoggerFactory.Instance;
        var taskStore = serviceProvider.GetKeyedService<ITaskStore>(agent.Name) ?? new InMemoryTaskStore();

        return new A2AServer(
            agentHandler,
            taskStore,
            new ChannelEventNotifier(),
            loggerFactory.CreateLogger<A2AServer>(),
            options?.ServerOptions);
    }
}