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

using System;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.AI.Agents;
using Moq;

namespace Microsoft.Agents.Orchestration.UnitTest;

public class HandoffsTests
{
    [Fact]
    public void EmptyConstructorsCreateEmptyCollections()
    {
        AgentHandoffs agentHandoffs = [];
        Assert.Empty(agentHandoffs);

        OrchestrationHandoffs orchestrationHandoffs = new("first");
        Assert.Empty(orchestrationHandoffs);
        Assert.Equal("first", orchestrationHandoffs.FirstAgentName);
    }

    [Fact]
    public void DictionaryConstructorsInvalidFirstAgent()
    {
        Assert.Throws<ArgumentNullException>(() => new OrchestrationHandoffs((string)null!));
        Assert.Throws<ArgumentException>(() => new OrchestrationHandoffs(string.Empty));
        Assert.Throws<ArgumentException>(() => new OrchestrationHandoffs(" "));
    }

    [Fact]
    public void AddWithAgentObjectsCreatesHandoffRelationships()
    {
        // Arrange
        OrchestrationHandoffs handoffs = new("source");

        AIAgent sourceAgent = CreateAgent("source", "Source Agent");
        AIAgent targetAgent1 = CreateAgent("target1", "Target Agent 1");
        AIAgent targetAgent2 = CreateAgent("target2", "Target Agent 2");

        // Act
        handoffs.Add(sourceAgent, targetAgent1, targetAgent2);

        // Assert
        Assert.Single(handoffs);
        Assert.Equal("source", handoffs.FirstAgentName);
        Assert.True(handoffs.ContainsKey("source"));

        AgentHandoffs sourceHandoffs = handoffs["source"];
        Assert.Equal(2, sourceHandoffs.Count);
        Assert.Equal("Target Agent 1", sourceHandoffs["target1"]);
        Assert.Equal("Target Agent 2", sourceHandoffs["target2"]);
    }

    [Fact]
    public void AddWithAgentAndCustomDescriptionUsesCustomDescription()
    {
        // Arrange
        OrchestrationHandoffs handoffs = new("source");

        AIAgent sourceAgent = CreateAgent("source", "Source Agent");
        AIAgent targetAgent = CreateAgent("target", "Target Agent");
        string customDescription = "Custom handoff description";

        // Act
        handoffs.Add(sourceAgent, targetAgent, customDescription);

        // Assert
        Assert.Single(handoffs);
        Assert.Equal("source", handoffs.FirstAgentName);
        AgentHandoffs sourceHandoffs = handoffs["source"];
        Assert.Single(sourceHandoffs);
        Assert.Equal(customDescription, sourceHandoffs["target"]);
    }

    [Fact]
    public void AddWithAgentAndTargetNameAddsHandoffWithDescription()
    {
        // Arrange
        OrchestrationHandoffs handoffs = new("source");

        AIAgent sourceAgent = CreateAgent("source", "Source Agent");
        string targetName = "targetName";
        string description = "Target description";

        // Act
        handoffs.Add(sourceAgent, targetName, description);

        // Assert
        Assert.Single(handoffs);
        Assert.Equal("source", handoffs.FirstAgentName);
        AgentHandoffs sourceHandoffs = handoffs["source"];
        Assert.Single(sourceHandoffs);
        Assert.Equal(description, sourceHandoffs[targetName]);
    }

    [Fact]
    public void AddWithSourceNameAndTargetNameAddsHandoffWithDescription()
    {
        // Arrange
        OrchestrationHandoffs handoffs = new("sourceName");

        string sourceName = "sourceName";
        string targetName = "targetName";
        string description = "Target description";

        // Act
        handoffs.Add(sourceName, targetName, description);

        // Assert
        Assert.Single(handoffs);
        Assert.Equal("sourceName", handoffs.FirstAgentName);
        AgentHandoffs sourceHandoffs = handoffs[sourceName];
        Assert.Single(sourceHandoffs);
        Assert.Equal(description, sourceHandoffs[targetName]);
    }

    [Fact]
    public void AddWithMultipleSourcesAndTargetsCreatesCorrectStructure()
    {
        // Arrange
        OrchestrationHandoffs handoffs = new("source1");

        AIAgent source1 = CreateAgent("source1", "Source Agent 1");
        AIAgent source2 = CreateAgent("source2", "Source Agent 2");

        AIAgent target1 = CreateAgent("target1", "Target Agent 1");
        AIAgent target2 = CreateAgent("target2", "Target Agent 2");
        AIAgent target3 = CreateAgent("target3", "Target Agent 3");

        // Act
        handoffs.Add(source1, target1, target2);
        handoffs.Add(source2, target2, target3);
        handoffs.Add(source1, target3, "Custom description");

        // Assert
        Assert.Equal(2, handoffs.Count);
        Assert.Equal("source1", handoffs.FirstAgentName);

        // Check source1's targets
        AgentHandoffs source1Handoffs = handoffs["source1"];
        Assert.Equal(3, source1Handoffs.Count);
        Assert.Equal("Target Agent 1", source1Handoffs["target1"]);
        Assert.Equal("Target Agent 2", source1Handoffs["target2"]);
        Assert.Equal("Custom description", source1Handoffs["target3"]);

        // Check source2's targets
        AgentHandoffs source2Handoffs = handoffs["source2"];
        Assert.Equal(2, source2Handoffs.Count);
        Assert.Equal("Target Agent 2", source2Handoffs["target2"]);
        Assert.Equal("Target Agent 3", source2Handoffs["target3"]);
    }

    [Fact]
    public void StaticAddCreatesNewOrchestrationHandoffs()
    {
        // Arrange
        AIAgent source = CreateAgent("source", "Source Agent");
        AIAgent target1 = CreateAgent("target1", "Target Agent 1");
        AIAgent target2 = CreateAgent("target2", "Target Agent 2");

        // Act
        OrchestrationHandoffs handoffs =
            OrchestrationHandoffs
                .StartWith(source)
                .Add(source, target1, target2);

        // Assert
        Assert.NotNull(handoffs);
        Assert.Equal(source.Id, handoffs.FirstAgentName);
        Assert.Single(handoffs);
        Assert.True(handoffs.ContainsKey("source"));

        AgentHandoffs sourceHandoffs = handoffs["source"];
        Assert.Equal(2, sourceHandoffs.Count);
        Assert.Equal("Target Agent 1", sourceHandoffs["target1"]);
        Assert.Equal("Target Agent 2", sourceHandoffs["target2"]);
    }

    [Fact]
    public void AddWithAgentsWithNoNameUsesId()
    {
        // Arrange
        OrchestrationHandoffs handoffs = new("source-id");

        AIAgent sourceAgent = CreateAgent(id: "source-id", name: null);
        AIAgent targetAgent = CreateAgent(id: "target-id", name: null, description: "Target Description");

        // Act
        handoffs.Add(sourceAgent, targetAgent);

        // Assert
        Assert.Single(handoffs);
        Assert.Equal("source-id", handoffs.FirstAgentName);
        Assert.True(handoffs.ContainsKey("source-id"));

        AgentHandoffs sourceHandoffs = handoffs["source-id"];
        Assert.Single(sourceHandoffs);
        Assert.Equal("Target Description", sourceHandoffs["target-id"]);
    }

    [Fact]
    public void AddWithAgentWithNoDescriptionUsesName()
    {
        // Arrange
        OrchestrationHandoffs handoffs = new("source");

        AIAgent sourceAgent = CreateAgent("source", "Source Agent");
        AIAgent targetAgent1 = CreateAgent("target1", name: "target 1");

        // Act
        handoffs.Add(sourceAgent, targetAgent1);

        // Assert
        Assert.Single(handoffs);
        Assert.Equal("source", handoffs.FirstAgentName);
        Assert.True(handoffs.ContainsKey("source"));

        AgentHandoffs sourceHandoffs = handoffs["source"];
        Assert.Single(sourceHandoffs);
        Assert.Equal("target 1", sourceHandoffs["target 1"]);
    }

    [Fact]
    public void AddWithAgentWithNoDescriptionOrNameThrows()
    {
        // Arrange
        OrchestrationHandoffs handoffs = new("source");

        AIAgent sourceAgent = CreateAgent("source", "Source Agent");
        AIAgent targetAgent1 = CreateAgent("target1");

        // Act
        InvalidOperationException ex = Assert.Throws<InvalidOperationException>(() => handoffs.Add(sourceAgent, targetAgent1));

        // Assert
        Assert.Equal("The provided target agent with Id 'target1' has no description or name, and no handoff description has been provided. At least one of these are required to register a handoff so that the appropriate target agent can be chosen.", ex.Message);
    }

    private static ChatClientAgent CreateAgent(string id, string? description = null, string? name = null)
    {
        Mock<IChatClient> mockClient = new(MockBehavior.Loose);
        ChatClientAgentOptions options =
            new()
            {
                Id = id,
                Name = name,
                Description = description,
            };
        ChatClientAgent mockAgent = new(mockClient.Object, options);
        return mockAgent;
    }
}