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

using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using FluentAssertions;
using Microsoft.Agents.AI.Workflows.Checkpointing;
using Microsoft.Agents.AI.Workflows.Sample;
using Microsoft.Agents.AI.Workflows.Specialized;
using Microsoft.Extensions.AI;

namespace Microsoft.Agents.AI.Workflows.UnitTests;

public class RepresentationTests
{
    private sealed class TestExecutor() : Executor("TestExecutor")
    {
        protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder protocolBuilder) => protocolBuilder;
    }

    private sealed class TestAgent : AIAgent
    {
        protected override ValueTask<AgentSession> CreateSessionCoreAsync(CancellationToken cancellationToken = default)
            => throw new NotImplementedException();

        protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
            => throw new NotImplementedException();

        protected override ValueTask<JsonElement> SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
            => throw new NotImplementedException();

        protected override Task<AgentResponse> RunCoreAsync(IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) =>
            throw new NotImplementedException();

        protected override IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) =>
            throw new NotImplementedException();
    }

    private static RequestPort TestRequestPort =>
        RequestPort.Create<FunctionCallContent, FunctionResultContent>("ExternalFunction");

    private static async ValueTask RunExecutorBindingInfoMatchTestAsync(ExecutorBinding binding)
    {
        ExecutorInfo info = binding.ToExecutorInfo();

        info.IsMatch(await binding.CreateInstanceAsync(sessionId: string.Empty)).Should().BeTrue();
    }

    [Fact]
    public async Task Test_ExecutorBinding_InfosAsync()
    {
        int testsRun = 0;
        await RunExecutorBindingTestAsync(new TestExecutor());
        await RunExecutorBindingTestAsync(TestRequestPort);
        await RunExecutorBindingTestAsync(new TestAgent());
        await RunExecutorBindingTestAsync(Step1EntryPoint.WorkflowInstance.BindAsExecutor(nameof(Step1EntryPoint)));

        Func<int, IWorkflowContext, CancellationToken, ValueTask> function = MessageHandlerAsync;
        await RunExecutorBindingTestAsync(function.BindAsExecutor("FunctionExecutor"));

        Type bindingBaseType = typeof(ExecutorBinding);
        Assembly workflowAssembly = bindingBaseType.Assembly;
        int expectedTests = workflowAssembly.GetTypes()
                                            .Count(type => type != bindingBaseType
                                                        && bindingBaseType.IsAssignableFrom(type));
        expectedTests.Should().BePositive();

        if (expectedTests > testsRun + 1)
        {
            Assert.Fail("Not all ExecutorBinding types were tested.");
        }

        async ValueTask RunExecutorBindingTestAsync(ExecutorBinding binding)
        {
            await RunExecutorBindingInfoMatchTestAsync(binding);
            testsRun++;
        }

        async ValueTask MessageHandlerAsync(int message, IWorkflowContext workflowContext, CancellationToken cancellationToken = default)
        {
        }
    }

    [Fact]
    public async Task Test_SpecializedExecutor_InfosAsync()
    {
        await RunExecutorBindingInfoMatchTestAsync(new AIAgentHostExecutor(new TestAgent(), new()));
        await RunExecutorBindingInfoMatchTestAsync(new RequestInfoExecutor(TestRequestPort));
    }

    private static string Source(int id) => $"Source/{id}";
    private static string Sink(int id) => $"Sink/{id}";

    private static Func<object?, bool> Condition() => Condition<object>();
    private static Func<TIn?, bool> Condition<TIn>() => _ => true;

    private static Func<object?, int, IEnumerable<int>> EdgeAssigner() => EdgeAssigner<object>();
    private static Func<TIn?, int, IEnumerable<int>> EdgeAssigner<TIn>() => (_, _) => [];

    [Fact]
    public void Test_EdgeInfos()
    {
        int edgeId = 0;

        // Direct Edges
        Edge directEdgeNoCondition = new(new DirectEdgeData(Source(1), Sink(2), TakeEdgeId()));
        RunEdgeInfoMatchTest(directEdgeNoCondition);

        Edge directEdgeNoCondition2 = new(new DirectEdgeData(Source(1), Sink(2), TakeEdgeId()));
        RunEdgeInfoMatchTest(directEdgeNoCondition, directEdgeNoCondition2);

        Edge directEdgeNoCondition3 = new(new DirectEdgeData(Source(3), Sink(4), TakeEdgeId()));
        RunEdgeInfoMatchTest(directEdgeNoCondition, directEdgeNoCondition3, expect: false);

        Edge directEdgeWithCondition = new(new DirectEdgeData(Source(3), Sink(4), TakeEdgeId(), Condition()));
        RunEdgeInfoMatchTest(directEdgeWithCondition);
        RunEdgeInfoMatchTest(directEdgeNoCondition2, directEdgeWithCondition, expect: false);
        RunEdgeInfoMatchTest(directEdgeNoCondition3, directEdgeWithCondition, expect: false);

        // FanOut Edges
        Edge fanOutEdgeNoAssigner = new(new FanOutEdgeData(Source(1), [Sink(2), Sink(3), Sink(4)], TakeEdgeId()));
        RunEdgeInfoMatchTest(fanOutEdgeNoAssigner);

        Edge fanOutEdgeNoAssigner2 = new(new FanOutEdgeData(Source(1), [Sink(2), Sink(3), Sink(4)], TakeEdgeId()));
        RunEdgeInfoMatchTest(fanOutEdgeNoAssigner, fanOutEdgeNoAssigner2);

        Edge fanOutEdgeNoAssigner3 = new(new FanOutEdgeData(Source(1), [Sink(3), Sink(4), Sink(2)], TakeEdgeId()));
        RunEdgeInfoMatchTest(fanOutEdgeNoAssigner, fanOutEdgeNoAssigner3, expect: false); // Order matters (though without Assigner maybe it shouldn't?)

        Edge fanOutEdgeNoAssigner4 = new(new FanOutEdgeData(Source(1), [Sink(2), Sink(3), Sink(5)], TakeEdgeId()));
        Edge fanOutEdgeNoAssigner5 = new(new FanOutEdgeData(Source(2), [Sink(2), Sink(3), Sink(4)], TakeEdgeId()));
        RunEdgeInfoMatchTest(fanOutEdgeNoAssigner, fanOutEdgeNoAssigner4, expect: false); // Identity matters
        RunEdgeInfoMatchTest(fanOutEdgeNoAssigner, fanOutEdgeNoAssigner5, expect: false);

        Edge fanOutEdgeWithAssigner = new(new FanOutEdgeData(Source(1), [Sink(2), Sink(3), Sink(4)], TakeEdgeId(), EdgeAssigner()));
        RunEdgeInfoMatchTest(fanOutEdgeWithAssigner);

        // FanIn Edges
        Edge fanInEdge = new(new FanInEdgeData([Source(1), Source(2), Source(3)], Sink(1), TakeEdgeId(), null));
        RunEdgeInfoMatchTest(fanInEdge);

        Edge fanInEdge2 = new(new FanInEdgeData([Source(1), Source(2), Source(3)], Sink(1), TakeEdgeId(), null));
        RunEdgeInfoMatchTest(fanInEdge, fanInEdge2);

        Edge fanInEdge3 = new(new FanInEdgeData([Source(2), Source(3), Source(1)], Sink(1), TakeEdgeId(), null));
        RunEdgeInfoMatchTest(fanInEdge, fanInEdge3, expect: false); // Order matters (though for FanIn maybe it shouldn't?)

        Edge fanInEdge4 = new(new FanInEdgeData([Source(1), Source(2), Source(4)], Sink(1), TakeEdgeId(), null));
        Edge fanInEdge5 = new(new FanInEdgeData([Source(1), Source(2), Source(3)], Sink(2), TakeEdgeId(), null));
        RunEdgeInfoMatchTest(fanInEdge, fanInEdge4, expect: false); // Identity matters
        RunEdgeInfoMatchTest(fanInEdge, fanInEdge5, expect: false);

        static void RunEdgeInfoMatchTest(Edge edge, Edge? comparatorEdge = null, bool expect = true)
        {
            comparatorEdge ??= edge;

            EdgeInfo info = edge.ToEdgeInfo();
            info.IsMatch(comparatorEdge).Should().Be(expect);
        }

        EdgeId TakeEdgeId() => new(edgeId++);
    }

    [Fact]
    public async Task Test_Sample_WorkflowInfosAsync()
    {
        RunWorkflowInfoMatchTest(Step1EntryPoint.WorkflowInstance);
        RunWorkflowInfoMatchTest(Step2EntryPoint.WorkflowInstance);
        RunWorkflowInfoMatchTest(Step3EntryPoint.WorkflowInstance);
        RunWorkflowInfoMatchTest(Step4EntryPoint.WorkflowInstance);
        // Step 5 reuses the workflow from Step 4, so we don't need to test it separately.
        RunWorkflowInfoMatchTest(Step6EntryPoint.CreateWorkflow(maxTurns: 2));
        // Step 7 reuses the workflow from Step 6, so we don't need to test it separately.

        RunWorkflowInfoMatchTest(Step1EntryPoint.WorkflowInstance, Step2EntryPoint.WorkflowInstance, expect: false);

        static void RunWorkflowInfoMatchTest(Workflow workflow, Workflow? comparator = null, bool expect = true)
        {
            comparator ??= workflow;

            WorkflowInfo info = workflow.ToWorkflowInfo();
            info.IsMatch(comparator).Should().Be(expect);
        }
    }
}