using System;
using System.Collections;
using System.Collections.Generic;
using System.IO;
using System.Threading;
using System.Threading.Tasks;

using AIRuntime.Runtime.Capabilities;
using AIRuntime.Runtime.Filesystem;
using AIRuntime.Runtime.Models;
using AIRuntime.Runtime.Providers;
using AIRuntime.Runtime.Providers.Models;
using AIRuntime.Runtime.Services.Agent;
using AIRuntime.Runtime.Capabilities.Permissions;

using NUnit.Framework;

using UnityEngine;
using UnityEngine.TestTools;


namespace AIRuntime.Tests
{
    public sealed class CapabilityCancellationTests
    {
        private string _workspace;


        [SetUp]
        public void Setup()
        {
            _workspace =
                Path.Combine(
                    Application.temporaryCachePath,
                    "AIRuntimeCapabilityCancellationTests");


            if (Directory.Exists(_workspace))
            {
                Directory.Delete(
                    _workspace,
                    true);
            }


            Directory.CreateDirectory(
                _workspace);
        }


        [TearDown]
        public void Cleanup()
        {
            if (Directory.Exists(_workspace))
            {
                Directory.Delete(
                    _workspace,
                    true);
            }
        }


        //--------------------------------------------------
        // Capability Receives Token
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator CapabilityReceivesCancellationToken()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                new AIRuntime.Runtime.AIRuntime(
                    _workspace);


            BlockingCapability capability =
                new BlockingCapability();


            runtime.RegisterCapability(
                capability);


            runtime.UseProvider(
                new BlockingCapabilityProvider());


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<string> task =
                runtime.ChatAsync(
                    "Run the blocking capability.",
                    cancellation.Token);


            while (!capability.HasStarted)
            {
                yield return null;
            }


            Assert.IsTrue(
                capability.ReceivedCancellableToken,
                "The capability should receive a cancellable token.");


            cancellation.Cancel();


            while (!task.IsCompleted)
            {
                yield return null;
            }


            Assert.IsTrue(
                task.IsCanceled);
        }


        //--------------------------------------------------
        // Cancellation Event
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator CancellationDuringCapability_EmitsRunCancelled()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                new AIRuntime.Runtime.AIRuntime(
                    _workspace);


            TestAgentObserver observer =
                new TestAgentObserver();


            BlockingCapability capability =
                new BlockingCapability();


            runtime.AddAgentObserver(
                observer);


            runtime.RegisterCapability(
                capability);


            runtime.UseProvider(
                new BlockingCapabilityProvider());


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<string> task =
                runtime.ChatAsync(
                    "Run and cancel the capability.",
                    cancellation.Token);


            while (!capability.HasStarted)
            {
                yield return null;
            }


            cancellation.Cancel();


            while (!task.IsCompleted)
            {
                yield return null;
            }


            Assert.IsTrue(
                task.IsCanceled);


            Assert.IsTrue(
                observer.Contains(
                    AgentEventType.RunStarted));


            Assert.IsTrue(
                observer.Contains(
                    AgentEventType.ToolStarted));


            Assert.IsTrue(
                observer.Contains(
                    AgentEventType.RunCancelled));


            Assert.IsFalse(
                observer.Contains(
                    AgentEventType.RunCompleted));


            Assert.IsFalse(
                observer.Contains(
                    AgentEventType.RunFailed));
        }


        //--------------------------------------------------
        // No ToolCompleted Event
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator CancellationDuringCapability_DoesNotEmitToolCompleted()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                new AIRuntime.Runtime.AIRuntime(
                    _workspace);


            TestAgentObserver observer =
                new TestAgentObserver();


            BlockingCapability capability =
                new BlockingCapability();


            runtime.AddAgentObserver(
                observer);


            runtime.RegisterCapability(
                capability);


            runtime.UseProvider(
                new BlockingCapabilityProvider());


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<string> task =
                runtime.ChatAsync(
                    "Cancel during tool execution.",
                    cancellation.Token);


            while (!capability.HasStarted)
            {
                yield return null;
            }


            cancellation.Cancel();


            while (!task.IsCompleted)
            {
                yield return null;
            }


            Assert.IsTrue(
                task.IsCanceled);


            Assert.IsTrue(
                observer.Contains(
                    AgentEventType.ToolStarted));


            Assert.IsFalse(
                observer.Contains(
                    AgentEventType.ToolCompleted));


            Assert.IsFalse(
                observer.Contains(
                    AgentEventType.ToolFailed));
        }


        //--------------------------------------------------
        // Tool Count Not Incremented
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator CancellationDuringCapability_DoesNotIncrementToolCount()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                new AIRuntime.Runtime.AIRuntime(
                    _workspace);


            BlockingCapability capability =
                new BlockingCapability();


            runtime.RegisterCapability(
                capability);


            runtime.UseProvider(
                new BlockingCapabilityProvider());


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<string> task =
                runtime.ChatAsync(
                    "Cancel the running capability.",
                    cancellation.Token);


            while (!capability.HasStarted)
            {
                yield return null;
            }


            cancellation.Cancel();


            while (!task.IsCompleted)
            {
                yield return null;
            }


            Assert.IsTrue(
                task.IsCanceled);


            Assert.IsNotNull(
                runtime.LastRunContext);


            Assert.AreEqual(
                0,
                runtime.LastRunContext.ToolCallsExecuted);
        }


        //--------------------------------------------------
        // Test Observer
        //--------------------------------------------------

        private sealed class TestAgentObserver :
            IAgentObserver
        {
            public List<AgentEvent> Events
            {
                get;
            } =
                new List<AgentEvent>();


            public void OnAgentEvent(
                AgentEvent agentEvent)
            {
                Events.Add(
                    agentEvent);
            }


            public bool Contains(
                AgentEventType eventType)
            {
                foreach (AgentEvent agentEvent in Events)
                {
                    if (agentEvent.Type == eventType)
                        return true;
                }


                return false;
            }
        }


        //--------------------------------------------------
        // Blocking Capability
        //--------------------------------------------------

        private sealed class BlockingCapability :
            CapabilityBase<BlockingArguments>
        {
            public bool HasStarted
            {
                get;
                private set;
            }

            public override CapabilityPermission Permission =>
                CapabilityPermission.ReadOnly;
            public bool ReceivedCancellableToken
            {
                get;
                private set;
            }


            public override string Name =>
                "blocking_capability";


            public override string Description =>
                "Blocks until the execution is cancelled.";


            protected override async Task<object> ExecuteAsync(
                BlockingArguments arguments,
                CapabilityContext context)
            {
                HasStarted =
                    true;


                ReceivedCancellableToken =
                    context.CancellationToken
                        .CanBeCanceled;


                await Task.Delay(
                    Timeout.Infinite,
                    context.CancellationToken);


                return "Unexpected completion.";
            }
        }


        private sealed class BlockingArguments
        {
        }


        //--------------------------------------------------
        // Provider Requesting Blocking Capability
        //--------------------------------------------------

        private sealed class BlockingCapabilityProvider :
            IProvider
        {
            private bool _toolRequested;


            public string Name =>
                "BlockingCapabilityProvider";


            public Task<ChatCompletion> ChatAsync(
                IReadOnlyList<ConversationMessage> messages,
                IReadOnlyList<ProviderTool> tools,
                CancellationToken cancellationToken = default)
            {
                cancellationToken
                    .ThrowIfCancellationRequested();


                if (_toolRequested)
                {
                    return Task.FromResult(
                        new ChatCompletion
                        {
                            Content =
                                "Unexpected final response."
                        });
                }


                _toolRequested =
                    true;


                ChatCompletion completion =
                    new ChatCompletion();


                completion.ToolCalls.Add(
                    new ToolCall
                    {
                        Id =
                            "blocking_capability_call",

                        Name =
                            "blocking_capability",

                        Arguments =
                            "{}"
                    });


                return Task.FromResult(
                    completion);
            }
        }
    }
}