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

using AIRuntime.Runtime.Models;
using AIRuntime.Runtime.Providers;
using AIRuntime.Runtime.Providers.Models;
using AIRuntime.Runtime.Services.Agent;

using NUnit.Framework;

using UnityEngine;
using UnityEngine.TestTools;


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


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


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


            Directory.CreateDirectory(
                _workspace);
        }


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


        //--------------------------------------------------
        // Agent Run Can Be Cancelled
        //--------------------------------------------------

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


            BlockingProvider provider =
                new BlockingProvider();


            runtime.UseProvider(
                provider);


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<string> task =
                runtime.ChatAsync(
                    "Wait until this run is cancelled.",
                    cancellation.Token);


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


            cancellation.Cancel();


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


            Assert.IsTrue(
                task.IsCanceled,
                "The agent task should be cancelled.");


            Assert.IsFalse(
                task.IsFaulted,
                "Cancellation should not be reported as a fault.");


            Assert.IsNotNull(
                runtime.LastRunContext);


            Assert.IsTrue(
                runtime.LastRunContext.CompletedAt.HasValue,
                "The cancelled run context should be completed.");
        }


        //--------------------------------------------------
        // Cancellation Event Is Emitted
        //--------------------------------------------------

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


            TestAgentObserver observer =
                new TestAgentObserver();


            BlockingProvider provider =
                new BlockingProvider();


            runtime.AddAgentObserver(
                observer);


            runtime.UseProvider(
                provider);


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<string> task =
                runtime.ChatAsync(
                    "Cancel this run.",
                    cancellation.Token);


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


            cancellation.Cancel();


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


            Assert.IsTrue(
                task.IsCanceled);


            Assert.AreEqual(
                2,
                observer.Events.Count);


            Assert.AreEqual(
                AgentEventType.RunStarted,
                observer.Events[0].Type);


            Assert.AreEqual(
                AgentEventType.RunCancelled,
                observer.Events[1].Type);


            Assert.AreEqual(
                observer.Events[0].SessionId,
                observer.Events[1].SessionId);


            Assert.AreEqual(
                runtime.LastRunContext.SessionId,
                observer.Events[1].SessionId);
        }


        //--------------------------------------------------
        // Cancelled Run Does Not Complete Normally
        //--------------------------------------------------

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


            TestAgentObserver observer =
                new TestAgentObserver();


            BlockingProvider provider =
                new BlockingProvider();


            runtime.AddAgentObserver(
                observer);


            runtime.UseProvider(
                provider);


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<string> task =
                runtime.ChatAsync(
                    "This run should be cancelled.",
                    cancellation.Token);


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


            cancellation.Cancel();


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


            Assert.IsTrue(
                task.IsCanceled);


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


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


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


        //--------------------------------------------------
        // Provider Failure Event Is Emitted
        //--------------------------------------------------

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


            TestAgentObserver observer =
                new TestAgentObserver();


            runtime.AddAgentObserver(
                observer);


            runtime.UseProvider(
                new FailingProvider());


            Task<string> task =
                runtime.ChatAsync(
                    "Trigger a provider failure.");


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


            Assert.IsTrue(
                task.IsFaulted,
                "The task should report the provider failure.");


            Assert.IsFalse(
                task.IsCanceled);


            Assert.AreEqual(
                2,
                observer.Events.Count);


            Assert.AreEqual(
                AgentEventType.RunStarted,
                observer.Events[0].Type);


            Assert.AreEqual(
                AgentEventType.RunFailed,
                observer.Events[1].Type);


            Assert.AreEqual(
                "Expected provider failure.",
                observer.Events[1].Message);


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


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


            Assert.IsNotNull(
                runtime.LastRunContext);


            Assert.IsTrue(
                runtime.LastRunContext.CompletedAt.HasValue);
        }


        //--------------------------------------------------
        // 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;
            }
        }


        //--------------------------------------------------
        // Provider That Waits Until Cancelled
        //--------------------------------------------------

        private sealed class BlockingProvider :
            IProvider
        {
            public string Name =>
                "BlockingProvider";


            public bool HasStarted
            {
                get;
                private set;
            }


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


                await Task.Delay(
                    Timeout.Infinite,
                    cancellationToken);


                throw new InvalidOperationException(
                    "The blocking provider should never complete normally.");
            }
        }


        //--------------------------------------------------
        // Provider That Always Fails
        //--------------------------------------------------

        private sealed class FailingProvider :
            IProvider
        {
            public string Name =>
                "FailingProvider";


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


                return Task.FromException<ChatCompletion>(
                    new InvalidOperationException(
                        "Expected provider failure."));
            }
        }
    }
}