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.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 RunTimeoutTests
    {
        private string _workspace;


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


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


            Directory.CreateDirectory(
                _workspace);
        }


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


        //--------------------------------------------------
        // Configuration Before Provider
        //--------------------------------------------------

        [Test]
        public void RunTimeout_CanBeConfiguredBeforeProvider()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                new AIRuntime.Runtime.AIRuntime(
                    _workspace);


            TimeSpan timeout =
                TimeSpan.FromSeconds(
                    12);


            runtime.RunTimeout =
                timeout;


            Assert.AreEqual(
                timeout,
                runtime.RunTimeout);


            Assert.DoesNotThrow(
                () =>
                {
                    runtime.UseProvider(
                        new SuccessfulProvider());
                });


            Assert.AreEqual(
                timeout,
                runtime.RunTimeout);
        }


        //--------------------------------------------------
        // Invalid Configuration
        //--------------------------------------------------

        [Test]
        public void RunTimeout_RejectsZero()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                new AIRuntime.Runtime.AIRuntime(
                    _workspace);


            ArgumentOutOfRangeException exception =
                Assert.Throws<ArgumentOutOfRangeException>(
                    () =>
                    {
                        runtime.RunTimeout =
                            TimeSpan.Zero;
                    });


            StringAssert.Contains(
                "Run timeout must be greater than zero or Timeout.InfiniteTimeSpan.",
                exception.Message);
        }


        [Test]
        public void RunTimeout_RejectsNegativeValue()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                new AIRuntime.Runtime.AIRuntime(
                    _workspace);


            ArgumentOutOfRangeException exception =
                Assert.Throws<ArgumentOutOfRangeException>(
                    () =>
                    {
                        runtime.RunTimeout =
                            TimeSpan.FromSeconds(
                                -2);
                    });


            StringAssert.Contains(
                "Run timeout must be greater than zero or Timeout.InfiniteTimeSpan.",
                exception.Message);
        }


        [Test]
        public void RunTimeout_AcceptsInfiniteTimeSpan()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                new AIRuntime.Runtime.AIRuntime(
                    _workspace);

            runtime.RunTimeout =
                Timeout.InfiniteTimeSpan;

            Assert.AreEqual(
                Timeout.InfiniteTimeSpan,
                runtime.RunTimeout);
        }


        //--------------------------------------------------
        // Provider Switching
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator RunTimeout_IsPreservedWhenProviderChanges()
        {
            SuccessfulProvider firstProvider =
                new SuccessfulProvider();


            SuccessfulProvider secondProvider =
                new SuccessfulProvider();


            AIRuntime.Runtime.AIRuntime runtime =
                new AIRuntime.Runtime.AIRuntime(
                    _workspace);


            runtime.RunTimeout =
                TimeSpan.FromSeconds(
                    8);


            runtime.UseProvider(
                firstProvider);


            Task<AgentRunResult> firstTask =
                runtime.RunAsync(
                    "First provider run.");


            yield return WaitForTask(
                firstTask);


            AssertTaskCompletedSuccessfully(
                firstTask,
                "First provider run");


            runtime.UseProvider(
                secondProvider);


            Assert.AreEqual(
                TimeSpan.FromSeconds(8),
                runtime.RunTimeout);


            Task<AgentRunResult> secondTask =
                runtime.RunAsync(
                    "Second provider run.");


            yield return WaitForTask(
                secondTask);


            AssertTaskCompletedSuccessfully(
                secondTask,
                "Second provider run");


            Assert.IsTrue(
                secondTask.Result.Succeeded);
        }


        //--------------------------------------------------
        // Full Run Timeout: Provider
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator RunTimeout_TimesOutBlockingProvider()
        {
            BlockingProvider provider =
                new BlockingProvider();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    provider);


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(
                    10);


            runtime.RunTimeout =
                TimeSpan.FromMilliseconds(
                    100);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Allow the full run to time out.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Full-run timeout");


            AgentRunResult result =
                task.GetAwaiter()
                    .GetResult();


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsTrue(
                result.TimedOut);


            StringAssert.Contains(
                "The agent run exceeded the maximum duration",
                result.ErrorMessage);


            Assert.AreEqual(
                0,
                runtime.MessageCount);


            Assert.IsFalse(
                runtime.IsRunning);
        }


        //--------------------------------------------------
        // Full Run Timeout: Capability
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator RunTimeout_TimesOutCapabilityExecution()
        {
            CapabilityRequestProvider provider =
                new CapabilityRequestProvider();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    provider);


            runtime.RegisterCapability(
                new BlockingCapability());


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(
                    10);


            runtime.RunTimeout =
                TimeSpan.FromMilliseconds(
                    100);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Execute the blocking capability.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Capability full-run timeout");


            AgentRunResult result =
                task.GetAwaiter()
                    .GetResult();


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsTrue(
                result.TimedOut);


            Assert.AreEqual(
                0,
                result.ToolCallsExecuted,
                "A cancelled capability should not be recorded as completed.");


            StringAssert.Contains(
                "The agent run exceeded the maximum duration",
                result.ErrorMessage);


            Assert.AreEqual(
                0,
                runtime.MessageCount);


            Assert.IsFalse(
                runtime.IsRunning);
        }


        //--------------------------------------------------
        // Persistent Conversation Rollback
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator RunTimeout_RollsBackCurrentPersistentRun()
        {
            SwitchableProvider provider =
                new SwitchableProvider();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    provider);


            Task<AgentRunResult> successfulTask =
                runtime.RunAsync(
                    "Persistent successful prompt.");


            yield return WaitForTask(
                successfulTask);


            AssertTaskCompletedSuccessfully(
                successfulTask,
                "Successful persistent run");


            Assert.IsTrue(
                successfulTask.Result.Succeeded);


            Assert.AreEqual(
                2,
                runtime.MessageCount);


            provider.ShouldBlock =
                true;


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(
                    10);


            runtime.RunTimeout =
                TimeSpan.FromMilliseconds(
                    100);


            Task<AgentRunResult> timedOutTask =
                runtime.RunAsync(
                    "This prompt must be rolled back.");


            yield return WaitForTask(
                timedOutTask);


            AssertTaskCompletedSuccessfully(
                timedOutTask,
                "Timed-out persistent run");


            AgentRunResult result =
                timedOutTask.GetAwaiter()
                    .GetResult();


            Assert.IsTrue(
                result.TimedOut);


            Assert.AreEqual(
                2,
                runtime.MessageCount,
                "The failed run should be rolled back without deleting earlier history.");
        }


        //--------------------------------------------------
        // Isolated Conversation
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator RunTimeout_IsolatedRunDoesNotModifyPersistentHistory()
        {
            SwitchableProvider provider =
                new SwitchableProvider();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    provider);


            Task<AgentRunResult> persistentTask =
                runtime.RunAsync(
                    "Store this conversation.");


            yield return WaitForTask(
                persistentTask);


            AssertTaskCompletedSuccessfully(
                persistentTask,
                "Persistent setup run");


            Assert.AreEqual(
                2,
                runtime.MessageCount);


            provider.ShouldBlock =
                true;


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(
                    10);


            runtime.RunTimeout =
                TimeSpan.FromMilliseconds(
                    100);


            Task<AgentRunResult> isolatedTask =
                runtime.RunIsolatedAsync(
                    "Time out this isolated task.");


            yield return WaitForTask(
                isolatedTask);


            AssertTaskCompletedSuccessfully(
                isolatedTask,
                "Timed-out isolated run");


            AgentRunResult result =
                isolatedTask.GetAwaiter()
                    .GetResult();


            Assert.IsTrue(
                result.TimedOut);


            Assert.IsFalse(
                result.Cancelled);


            Assert.AreEqual(
                2,
                runtime.MessageCount,
                "The isolated run must not alter persistent history.");
        }


        //--------------------------------------------------
        // Manual Cancellation Remains Distinct
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ManualCancellation_IsNotReportedAsTimeout()
        {
            BlockingProvider provider =
                new BlockingProvider();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    provider);


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(
                    10);


            runtime.RunTimeout =
                TimeSpan.FromSeconds(
                    10);


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Cancel this run manually.",
                    cancellation.Token);


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


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Manually cancelled structured run");


            AgentRunResult result =
                task.GetAwaiter()
                    .GetResult();


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsTrue(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


            Assert.AreEqual(
                "The agent run was cancelled.",
                result.ErrorMessage);


            Assert.IsFalse(
                runtime.IsRunning);
        }


        //--------------------------------------------------
        // Provider Timeout Remains Independent
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ProviderTimeout_StillWorksIndependently()
        {
            BlockingProvider provider =
                new BlockingProvider();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    provider);


            runtime.ProviderTimeout =
                TimeSpan.FromMilliseconds(
                    100);


            runtime.RunTimeout =
                TimeSpan.FromSeconds(
                    5);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Allow the provider request to time out.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Provider timeout run");


            AgentRunResult result =
                task.GetAwaiter()
                    .GetResult();


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsTrue(
                result.TimedOut);


            StringAssert.Contains(
                "Provider 'BlockingProvider' did not respond within",
                result.ErrorMessage);


            StringAssert.DoesNotContain(
                "agent run exceeded",
                result.ErrorMessage.ToLowerInvariant());
        }


        //--------------------------------------------------
        // Runtime Recovery
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator Runtime_IsAvailableAfterRunTimeout()
        {
            SwitchableProvider provider =
                new SwitchableProvider
                {
                    ShouldBlock =
                        true
                };


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    provider);


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(
                    10);


            runtime.RunTimeout =
                TimeSpan.FromMilliseconds(
                    100);


            Task<AgentRunResult> timeoutTask =
                runtime.RunAsync(
                    "Time out this run.");


            yield return WaitForTask(
                timeoutTask);


            AssertTaskCompletedSuccessfully(
                timeoutTask,
                "Initial timed-out run");


            Assert.IsTrue(
                timeoutTask.Result.TimedOut);


            Assert.IsFalse(
                runtime.IsRunning);


            provider.ShouldBlock =
                false;


            runtime.RunTimeout =
                TimeSpan.FromSeconds(
                    5);


            Task<AgentRunResult> recoveryTask =
                runtime.RunAsync(
                    "Run successfully afterward.");


            yield return WaitForTask(
                recoveryTask);


            AssertTaskCompletedSuccessfully(
                recoveryTask,
                "Recovery run");


            AgentRunResult recoveryResult =
                recoveryTask.GetAwaiter()
                    .GetResult();


            Assert.IsTrue(
                recoveryResult.Succeeded);


            Assert.IsFalse(
                recoveryResult.Cancelled);


            Assert.IsFalse(
                recoveryResult.TimedOut);


            Assert.AreEqual(
                "Successful response.",
                recoveryResult.Response);
        }


        //--------------------------------------------------
        // Configuration During Active Run
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ChangingRunTimeoutDuringRun_IsRejected()
        {
            BlockingProvider provider =
                new BlockingProvider();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    provider);


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(
                    10);


            runtime.RunTimeout =
                TimeSpan.FromSeconds(
                    10);


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Keep this run active.",
                    cancellation.Token);


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


            InvalidOperationException exception =
                Assert.Throws<InvalidOperationException>(
                    () =>
                    {
                        runtime.RunTimeout =
                            TimeSpan.FromSeconds(
                                20);
                    });


            Assert.AreEqual(
                "The run timeout cannot be changed while an agent run is in progress.",
                exception.Message);


            Assert.AreEqual(
                TimeSpan.FromSeconds(10),
                runtime.RunTimeout);


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Cancelled cleanup run");


            Assert.IsTrue(
                task.Result.Cancelled);
        }


        //--------------------------------------------------
        // Timeout Event
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator RunTimeout_EmitsSingleTimedOutEvent()
        {
            BlockingProvider provider =
                new BlockingProvider();


            TestObserver observer =
                new TestObserver();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    provider);


            runtime.AddAgentObserver(
                observer);


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(
                    10);


            runtime.RunTimeout =
                TimeSpan.FromMilliseconds(
                    100);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Observe the timeout.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Observed timeout run");


            Assert.IsTrue(
                task.Result.TimedOut);


            Assert.AreEqual(
                1,
                observer.Count(
                    AgentEventType.RunStarted));


            Assert.AreEqual(
                1,
                observer.Count(
                    AgentEventType.RunTimedOut));


            Assert.AreEqual(
                0,
                observer.Count(
                    AgentEventType.RunCancelled));


            Assert.AreEqual(
                0,
                observer.Count(
                    AgentEventType.RunFailed));


            Assert.AreEqual(
                0,
                observer.Count(
                    AgentEventType.RunCompleted));
        }


        //--------------------------------------------------
        // Runtime Creation
        //--------------------------------------------------

        private AIRuntime.Runtime.AIRuntime CreateRuntime(
            IProvider provider)
        {
            AIRuntime.Runtime.AIRuntime runtime =
                new AIRuntime.Runtime.AIRuntime(
                    _workspace);


            runtime.UseProvider(
                provider);


            return runtime;
        }


        //--------------------------------------------------
        // Capabilities
        //--------------------------------------------------

        private sealed class BlockingCapability :
            CapabilityBase<BlockingArguments>
        {
            public override string Name =>
                "blocking_capability";

            public override CapabilityPermission Permission =>
                CapabilityPermission.ReadOnly;

            public override string Description =>
                "Waits indefinitely until execution is cancelled.";


            protected override async Task<object> ExecuteAsync(
                BlockingArguments arguments,
                CapabilityContext context)
            {
                await Task.Delay(
                    Timeout.Infinite,
                    context.CancellationToken);


                return "Unexpected completion.";
            }
        }


        private sealed class BlockingArguments
        {
        }


        //--------------------------------------------------
        // Providers
        //--------------------------------------------------

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


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


                return Task.FromResult(
                    new ChatCompletion
                    {
                        Content =
                            "Successful response."
                    });
            }
        }


        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.");
            }
        }


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


            public bool ShouldBlock
            {
                get;
                set;
            }


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


                if (!ShouldBlock)
                {
                    return new ChatCompletion
                    {
                        Content =
                            "Successful response."
                    };
                }


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


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


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


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


                ChatCompletion completion =
                    new ChatCompletion();


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

                        Name =
                            "blocking_capability",

                        Arguments =
                            "{}"
                    });


                return Task.FromResult(
                    completion);
            }
        }


        //--------------------------------------------------
        // Observer
        //--------------------------------------------------

        private sealed class TestObserver :
            IAgentObserver
        {
            private readonly List<AgentEvent> _events =
                new List<AgentEvent>();


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


            public int Count(
                AgentEventType type)
            {
                int count =
                    0;


                foreach (AgentEvent agentEvent in _events)
                {
                    if (agentEvent.Type == type)
                    {
                        count++;
                    }
                }


                return count;
            }
        }


        //--------------------------------------------------
        // Coroutine Helpers
        //--------------------------------------------------

        private static IEnumerator WaitForTask(
            Task task)
        {
            while (!task.IsCompleted)
            {
                yield return null;
            }
        }


        //--------------------------------------------------
        // Assertion Helpers
        //--------------------------------------------------

        private static void AssertTaskCompletedSuccessfully(
            Task task,
            string operation)
        {
            Assert.IsNotNull(
                task,
                operation + " returned a null task.");


            if (task.IsCanceled)
            {
                Assert.Fail(
                    operation +
                    " was unexpectedly cancelled.");
            }


            if (!task.IsFaulted)
                return;


            Exception exception =
                UnwrapException(
                    task.Exception);


            Assert.Fail(
                operation +
                " failed:\n" +
                (exception?.ToString() ??
                 "Unknown task failure."));
        }


        private static Exception UnwrapException(
            Exception exception)
        {
            if (!(exception is AggregateException aggregate))
                return exception;


            AggregateException flattened =
                aggregate.Flatten();


            if (flattened.InnerExceptions.Count == 1)
            {
                return flattened.InnerExceptions[0];
            }


            return flattened;
        }
    }
}
