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

using System.ComponentModel;
using System.Diagnostics;
using System.Reflection;
using Microsoft.Agents.AI.DurableTask.IntegrationTests.Logging;
using Microsoft.DurableTask;
using Microsoft.DurableTask.Client;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.Configuration;
using OpenAI.Chat;

namespace Microsoft.Agents.AI.DurableTask.IntegrationTests;

/// <summary>
/// Tests for scenarios where an external client interacts with Durable Task Agents.
/// </summary>
[Collection("Sequential")]
[Trait("Category", "Integration")]
public sealed class ExternalClientTests(ITestOutputHelper outputHelper) : IDisposable
{
    private static readonly TimeSpan s_defaultTimeout = Debugger.IsAttached
        ? TimeSpan.FromMinutes(5)
        : TimeSpan.FromSeconds(120);

    private static readonly IConfiguration s_configuration =
        new ConfigurationBuilder()
            .AddUserSecrets(Assembly.GetExecutingAssembly())
            .AddEnvironmentVariables()
            .Build();

    private readonly ITestOutputHelper _outputHelper = outputHelper;
    private readonly CancellationTokenSource _cts = new(delay: s_defaultTimeout);

    private CancellationToken TestTimeoutToken => this._cts.Token;

    public void Dispose() => this._cts.Dispose();

    [RetryFact(2, 5000)]
    public async Task SimplePromptAsync()
    {
        // Setup
        AIAgent simpleAgent = TestHelper.GetAzureOpenAIChatClient(s_configuration).AsAIAgent(
            instructions: "You are a helpful assistant that always responds with a friendly greeting.",
            name: "TestAgent");

        using TestHelper testHelper = TestHelper.Start([simpleAgent], this._outputHelper);

        // A proxy agent is needed to call the hosted test agent
        AIAgent simpleAgentProxy = simpleAgent.AsDurableAgentProxy(testHelper.Services);

        // Act: send a prompt to the agent and wait for a response
        AgentSession session = await simpleAgentProxy.CreateSessionAsync(this.TestTimeoutToken);
        await simpleAgentProxy.RunAsync(
            message: "Hello!",
            session,
            cancellationToken: this.TestTimeoutToken);

        AgentResponse response = await simpleAgentProxy.RunAsync(
            message: "Repeat what you just said but say it like a pirate",
            session,
            cancellationToken: this.TestTimeoutToken);

        // Assert: verify the agent responded appropriately
        // We can't predict the exact response, but we can check that there is one response
        Assert.NotNull(response);
        Assert.NotEmpty(response.Text);

        // Assert: verify the expected log entries were created in the expected category
        IReadOnlyCollection<LogEntry> logs = testHelper.GetLogs();
        Assert.NotEmpty(logs);
        List<LogEntry> agentLogs = [.. logs.Where(log => log.Category.Contains(simpleAgent.Name!)).ToList()];
        Assert.NotEmpty(agentLogs);
        Assert.Contains(agentLogs, log => log.EventId.Name == "LogAgentRequest" && log.Message.Contains("Hello!"));
        Assert.Contains(agentLogs, log => log.EventId.Name == "LogAgentResponse");
    }

    [RetryFact(2, 5000)]
    public async Task CallFunctionToolsAsync()
    {
        int weatherToolInvocationCount = 0;
        int packingListToolInvocationCount = 0;

        string GetWeather(string location)
        {
            weatherToolInvocationCount++;
            return $"The weather in {location} is sunny with a high of 75°F and a low of 55°F.";
        }

        string SuggestPackingList(string weather, bool isSunny)
        {
            packingListToolInvocationCount++;
            return isSunny ? "Pack sunglasses and sunscreen." : "Pack a raincoat and umbrella.";
        }

        AIAgent tripPlanningAgent = TestHelper.GetAzureOpenAIChatClient(s_configuration).AsAIAgent(
            instructions: "You are a trip planning assistant. Use the weather tool and packing list tool as needed.",
            name: "TripPlanningAgent",
            description: "An agent to help plan your day trips",
            tools: [AIFunctionFactory.Create(GetWeather), AIFunctionFactory.Create(SuggestPackingList)]
        );

        using TestHelper testHelper = TestHelper.Start([tripPlanningAgent], this._outputHelper);
        AIAgent tripPlanningAgentProxy = tripPlanningAgent.AsDurableAgentProxy(testHelper.Services);

        // Act: send a prompt to the agent
        AgentResponse response = await tripPlanningAgentProxy.RunAsync(
            message: "Help me figure out what to pack for my Seattle trip next Sunday",
            cancellationToken: this.TestTimeoutToken);

        // Assert: verify the agent responded appropriately
        // We can't predict the exact response, but we can check that there is one response
        Assert.NotNull(response);
        Assert.NotEmpty(response.Text);

        // Assert: verify the expected log entries were created in the expected category
        IReadOnlyCollection<LogEntry> logs = testHelper.GetLogs();
        Assert.NotEmpty(logs);

        List<LogEntry> agentLogs = [.. logs.Where(log => log.Category.Contains(tripPlanningAgent.Name!)).ToList()];
        Assert.NotEmpty(agentLogs);
        Assert.Contains(agentLogs, log => log.EventId.Name == "LogAgentRequest" && log.Message.Contains("Seattle trip"));
        Assert.Contains(agentLogs, log => log.EventId.Name == "LogAgentResponse");

        // Assert: verify the tools were called
        Assert.Equal(1, weatherToolInvocationCount);
        Assert.Equal(1, packingListToolInvocationCount);
    }

    [RetryFact(2, 5000)]
    public async Task CallLongRunningFunctionToolsAsync()
    {
        [Description("Starts a greeting workflow and returns the workflow instance ID")]
        string StartWorkflowTool(string name)
        {
            return DurableAgentContext.Current.ScheduleNewOrchestration(nameof(RunWorkflowAsync), input: name);
        }

        [Description("Gets the current status of a previously started workflow. A null response means the workflow has not started yet.")]
        static async Task<OrchestrationMetadata?> GetWorkflowStatusToolAsync(string instanceId)
        {
            OrchestrationMetadata? status = await DurableAgentContext.Current.GetOrchestrationStatusAsync(
                instanceId,
                includeDetails: true);
            if (status == null)
            {
                // If the status is not found, wait a bit before returning null to give the workflow time to start
                await Task.Delay(TimeSpan.FromSeconds(1));
            }

            return status;
        }

        async Task<string> RunWorkflowAsync(TaskOrchestrationContext context, string name)
        {
            // 1. Get agent and create a session
            DurableAIAgent agent = context.GetAgent("SimpleAgent");
            AgentSession session = await agent.CreateSessionAsync(this.TestTimeoutToken);

            // 2. Call an agent and tell it my name
            await agent.RunAsync($"My name is {name}.", session);

            // 3. Call the agent again with the same session (ask it to tell me my name)
            AgentResponse response = await agent.RunAsync("What is my name?", session);

            return response.Text;
        }

        using TestHelper testHelper = TestHelper.Start(
            this._outputHelper,
            configureAgents: agents =>
            {
                // This is the agent that will be used to start the workflow
                agents.AddAIAgentFactory(
                    "WorkflowAgent",
                    sp => TestHelper.GetAzureOpenAIChatClient(s_configuration).AsAIAgent(
                        name: "WorkflowAgent",
                        instructions: "You can start greeting workflows and check their status.",
                        services: sp,
                        tools: [
                            AIFunctionFactory.Create(StartWorkflowTool),
                            AIFunctionFactory.Create(GetWorkflowStatusToolAsync)
                        ]));

                // This is the agent that will be called by the workflow
                agents.AddAIAgent(TestHelper.GetAzureOpenAIChatClient(s_configuration).AsAIAgent(
                    name: "SimpleAgent",
                    instructions: "You are a simple assistant."
                ));
            },
            durableTaskRegistry: registry => registry.AddOrchestratorFunc<string, string>(nameof(RunWorkflowAsync), RunWorkflowAsync));

        AIAgent workflowManagerAgentProxy = testHelper.Services.GetDurableAgentProxy("WorkflowAgent");

        // Act: send a prompt to the agent
        AgentSession session = await workflowManagerAgentProxy.CreateSessionAsync(this.TestTimeoutToken);
        await workflowManagerAgentProxy.RunAsync(
            message: "Start a greeting workflow for \"John Doe\".",
            session,
            cancellationToken: this.TestTimeoutToken);

        // Act: prompt it again to wait for the workflow to complete
        AgentResponse response = await workflowManagerAgentProxy.RunAsync(
            message: "Wait for the workflow to complete and tell me the result.",
            session,
            cancellationToken: this.TestTimeoutToken);

        // Assert: verify the agent responded appropriately
        // We can't predict the exact response, but we can check that there is one response
        Assert.NotNull(response);
        Assert.NotEmpty(response.Text);
        Assert.Contains("John Doe", response.Text);
    }

    [Fact]
    public void AsDurableAgentProxy_ThrowsWhenAgentNotRegistered()
    {
        // Setup: Register one agent but try to use a different one
        AIAgent registeredAgent = TestHelper.GetAzureOpenAIChatClient(s_configuration).AsAIAgent(
            instructions: "You are a helpful assistant.",
            name: "RegisteredAgent");

        using TestHelper testHelper = TestHelper.Start([registeredAgent], this._outputHelper);

        // Create an agent with a different name that isn't registered
        AIAgent unregisteredAgent = TestHelper.GetAzureOpenAIChatClient(s_configuration).AsAIAgent(
            instructions: "You are a helpful assistant.",
            name: "UnregisteredAgent");

        // Act & Assert: Should throw AgentNotRegisteredException
        AgentNotRegisteredException exception = Assert.Throws<AgentNotRegisteredException>(
            () => unregisteredAgent.AsDurableAgentProxy(testHelper.Services));

        Assert.Equal("UnregisteredAgent", exception.AgentName);
    }
}