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


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


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


            Directory.CreateDirectory(
                _workspace);
        }


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


        //--------------------------------------------------
        // Convenience Property Mapping
        //--------------------------------------------------

        [Test]
        public void SucceededStatus_ActivatesOnlySucceededProperty()
        {
            AgentRunResult result =
                new AgentRunResult
                {
                    Status =
                        AgentRunStatus.Succeeded
                };


            Assert.AreEqual(
                AgentRunStatus.Succeeded,
                result.Status);


            Assert.IsTrue(
                result.Succeeded);


            Assert.IsFalse(
                result.Failed);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);
        }


        [Test]
        public void FailedStatus_ActivatesOnlyFailedProperty()
        {
            AgentRunResult result =
                new AgentRunResult
                {
                    Status =
                        AgentRunStatus.Failed
                };


            Assert.AreEqual(
                AgentRunStatus.Failed,
                result.Status);


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsTrue(
                result.Failed);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);
        }


        [Test]
        public void CancelledStatus_ActivatesOnlyCancelledProperty()
        {
            AgentRunResult result =
                new AgentRunResult
                {
                    Status =
                        AgentRunStatus.Cancelled
                };


            Assert.AreEqual(
                AgentRunStatus.Cancelled,
                result.Status);


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Failed);


            Assert.IsTrue(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);
        }


        [Test]
        public void TimedOutStatus_ActivatesOnlyTimedOutProperty()
        {
            AgentRunResult result =
                new AgentRunResult
                {
                    Status =
                        AgentRunStatus.TimedOut
                };


            Assert.AreEqual(
                AgentRunStatus.TimedOut,
                result.Status);


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Failed);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsTrue(
                result.TimedOut);
        }


        //--------------------------------------------------
        // Successful Execution
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator SuccessfulRun_ReturnsSucceededStatus()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    new SuccessfulProvider());


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Complete successfully.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Successful structured run");


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


            Assert.AreEqual(
                AgentRunStatus.Succeeded,
                result.Status);


            Assert.IsTrue(
                result.Succeeded);


            Assert.IsFalse(
                result.Failed);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


            Assert.AreEqual(
                "Successful response.",
                result.Response);


            Assert.IsTrue(
                string.IsNullOrEmpty(
                    result.ErrorMessage));
        }


        //--------------------------------------------------
        // Ordinary Failure
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ProviderFailure_ReturnsFailedStatus()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    new FailingProvider());


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


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Failed structured run");


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


            Assert.AreEqual(
                AgentRunStatus.Failed,
                result.Status);


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsTrue(
                result.Failed);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


            Assert.AreEqual(
                string.Empty,
                result.Response);


            Assert.AreEqual(
                "Simulated provider failure.",
                result.ErrorMessage);


            Assert.AreEqual(
                0,
                runtime.MessageCount);
        }


        //--------------------------------------------------
        // Manual Cancellation
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ManualCancellation_ReturnsCancelledStatus()
        {
            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.",
                    cancellation.Token);


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


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Cancelled structured run");


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


            Assert.AreEqual(
                AgentRunStatus.Cancelled,
                result.Status);


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Failed);


            Assert.IsTrue(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


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


            Assert.AreEqual(
                0,
                runtime.MessageCount);


            Assert.IsFalse(
                runtime.IsRunning);
        }


        //--------------------------------------------------
        // Provider Timeout
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ProviderTimeout_ReturnsTimedOutStatus()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    new BlockingProvider());


            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 structured run");


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


            Assert.AreEqual(
                AgentRunStatus.TimedOut,
                result.Status);


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Failed);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsTrue(
                result.TimedOut);


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


            Assert.AreEqual(
                0,
                runtime.MessageCount);


            Assert.IsFalse(
                runtime.IsRunning);
        }


        //--------------------------------------------------
        // Full Run Timeout
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator FullRunTimeout_ReturnsTimedOutStatus()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    new BlockingProvider());


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(
                    10);


            runtime.RunTimeout =
                TimeSpan.FromMilliseconds(
                    100);


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


            yield return WaitForTask(
                task);


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


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


            Assert.AreEqual(
                AgentRunStatus.TimedOut,
                result.Status);


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Failed);


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


        //--------------------------------------------------
        // Capability Failure
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator CapabilityFailure_ReturnsFailedStatus()
        {
            CapabilityFailureProvider provider =
                new CapabilityFailureProvider();


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


            runtime.RegisterCapability(
                new FailingCapability());


            LogAssert.Expect(
                LogType.Error,
                "[Capability] 'failing_capability' failed: Simulated capability failure.");


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


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Capability-failure structured run");


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


            Assert.AreEqual(
                AgentRunStatus.Failed,
                result.Status);


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsTrue(
                result.Failed);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


            Assert.AreEqual(
                0,
                result.ToolCallsExecuted);


            Assert.AreEqual(
                "Simulated capability failure.",
                result.ErrorMessage);


            Assert.AreEqual(
                0,
                runtime.MessageCount);
        }


        //--------------------------------------------------
        // Capability Round Limit
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator CapabilityRoundLimit_ReturnsFailedStatus()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    new RepeatingCapabilityProvider());


            runtime.RegisterCapability(
                new EchoCapability());


            runtime.MaxCapabilityRounds =
                2;


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Continue calling capabilities.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Capability-round-limit run");


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


            Assert.AreEqual(
                AgentRunStatus.Failed,
                result.Status);


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsTrue(
                result.Failed);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


            Assert.AreEqual(
                2,
                result.ToolCallsExecuted);


            Assert.AreEqual(
                "Conversation exceeded the maximum of 2 capability rounds.",
                result.ErrorMessage);


            Assert.AreEqual(
                0,
                runtime.MessageCount);
        }


        //--------------------------------------------------
        // Isolated Execution
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator SuccessfulIsolatedRun_ReturnsSucceededStatus()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    new SuccessfulProvider());


            Task<AgentRunResult> task =
                runtime.RunIsolatedAsync(
                    "Complete this isolated run.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Successful isolated run");


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


            Assert.AreEqual(
                AgentRunStatus.Succeeded,
                result.Status);


            Assert.IsTrue(
                result.Succeeded);


            Assert.IsFalse(
                result.Failed);


            Assert.AreEqual(
                0,
                runtime.MessageCount);
        }


        [UnityTest]
        public IEnumerator FailedIsolatedRun_ReturnsFailedStatus()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    new FailingProvider());


            Task<AgentRunResult> task =
                runtime.RunIsolatedAsync(
                    "Fail this isolated run.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Failed isolated run");


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


            Assert.AreEqual(
                AgentRunStatus.Failed,
                result.Status);


            Assert.IsTrue(
                result.Failed);


            Assert.IsFalse(
                result.Succeeded);


            Assert.AreEqual(
                0,
                runtime.MessageCount);
        }


        //--------------------------------------------------
        // Status Can Change Without Contradiction
        //--------------------------------------------------

        [Test]
        public void ChangingStatus_UpdatesAllConvenienceProperties()
        {
            AgentRunResult result =
                new AgentRunResult();


            result.Status =
                AgentRunStatus.Succeeded;


            Assert.IsTrue(
                result.Succeeded);


            result.Status =
                AgentRunStatus.Failed;


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsTrue(
                result.Failed);


            result.Status =
                AgentRunStatus.Cancelled;


            Assert.IsFalse(
                result.Failed);


            Assert.IsTrue(
                result.Cancelled);


            result.Status =
                AgentRunStatus.TimedOut;


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsTrue(
                result.TimedOut);
        }


        //--------------------------------------------------
        // 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 EchoCapability :
            CapabilityBase<EchoArguments>
        {
            public override string Name =>
                "echo";

            public override CapabilityPermission Permission =>
                CapabilityPermission.ReadOnly;
            public override string Description =>
                "Returns the supplied text.";


            protected override Task<object> ExecuteAsync(
                EchoArguments arguments,
                CapabilityContext context)
            {
                context.CancellationToken
                    .ThrowIfCancellationRequested();


                return Task.FromResult<object>(
                    arguments.Text ??
                    string.Empty);
            }
        }


        private sealed class EchoArguments
        {
            public string Text
            {
                get;
                set;
            }
        }


        private sealed class FailingCapability :
            CapabilityBase<FailingArguments>
        {
            public override string Name =>
                "failing_capability";

            public override CapabilityPermission Permission =>
                CapabilityPermission.ReadOnly;

            public override string Description =>
                "Always throws a test exception.";


            protected override Task<object> ExecuteAsync(
                FailingArguments arguments,
                CapabilityContext context)
            {
                context.CancellationToken
                    .ThrowIfCancellationRequested();


                // InvalidOperationException is intentionally recoverable for
                // capability calls whose arguments or current Editor state
                // can be corrected by the model. This fixture represents a
                // terminal capability failure, so use an unclassified
                // exception and verify the run fails immediately.
                throw new Exception(
                    "Simulated capability failure.");
            }
        }


        private sealed class FailingArguments
        {
        }


        //--------------------------------------------------
        // 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 FailingProvider :
            IProvider
        {
            public string Name =>
                "FailingProvider";


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


                throw new InvalidOperationException(
                    "Simulated provider failure.");
            }
        }


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


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


            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 =
                            "failing_call_1",

                        Name =
                            "failing_capability",

                        Arguments =
                            "{}"
                    });


                return Task.FromResult(
                    completion);
            }
        }


        private sealed class RepeatingCapabilityProvider :
            IProvider
        {
            private int _requestNumber;


            public string Name =>
                "RepeatingCapabilityProvider";


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


                _requestNumber++;


                ChatCompletion completion =
                    new ChatCompletion();


                completion.ToolCalls.Add(
                    new ToolCall
                    {
                        Id =
                            "echo_call_" +
                            _requestNumber,

                        Name =
                            "echo",

                        Arguments =
                            "{\"Text\":\"round-" +
                            _requestNumber +
                            "\"}"
                    });


                return Task.FromResult(
                    completion);
            }
        }


        //--------------------------------------------------
        // Coroutine Helper
        //--------------------------------------------------

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