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.Capabilities.Permissions;
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 IsolatedRunTests
    {
        private string _workspace;


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


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


            Directory.CreateDirectory(
                _workspace);
        }


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


        //--------------------------------------------------
        // Existing History Is Not Supplied
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator IsolatedRun_DoesNotReceiveExistingHistory()
        {
            RecordingProvider provider =
                new RecordingProvider();


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


            Task<string> normalTask =
                runtime.ChatAsync(
                    "Persistent prompt.");


            yield return WaitForTask(
                normalTask);


            AssertTaskCompletedSuccessfully(
                normalTask,
                "Persistent conversation run");


            Assert.AreEqual(
                2,
                runtime.MessageCount);


            Task<AgentRunResult> isolatedTask =
                runtime.RunIsolatedAsync(
                    "Isolated prompt.");


            yield return WaitForTask(
                isolatedTask);


            AssertTaskCompletedSuccessfully(
                isolatedTask,
                "Isolated run");


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


            Assert.IsTrue(
                result.Succeeded);


            Assert.AreEqual(
                2,
                provider.RequestSnapshots.Count);


            IReadOnlyList<ConversationMessage> isolatedRequest =
                provider.RequestSnapshots[1];


            Assert.AreEqual(
                1,
                isolatedRequest.Count,
                "An isolated run should receive only its own user prompt.");


            Assert.AreEqual(
                "user",
                isolatedRequest[0].Role);


            Assert.AreEqual(
                "Isolated prompt.",
                isolatedRequest[0].Content);
        }


        //--------------------------------------------------
        // Persistent History Is Not Modified
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator IsolatedRun_DoesNotModifyExistingHistory()
        {
            RecordingProvider provider =
                new RecordingProvider();


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


            Task<string> normalTask =
                runtime.ChatAsync(
                    "Keep this conversation.");


            yield return WaitForTask(
                normalTask);


            AssertTaskCompletedSuccessfully(
                normalTask,
                "Persistent conversation run");


            Assert.AreEqual(
                2,
                runtime.MessageCount);


            Task<AgentRunResult> isolatedTask =
                runtime.RunIsolatedAsync(
                    "Do not store this task.");


            yield return WaitForTask(
                isolatedTask);


            AssertTaskCompletedSuccessfully(
                isolatedTask,
                "Isolated run");


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


            Assert.IsTrue(
                result.Succeeded);


            Assert.AreEqual(
                2,
                runtime.MessageCount,
                "The isolated run must not add messages to persistent history.");


            Task<string> nextNormalTask =
                runtime.ChatAsync(
                    "Continue the persistent conversation.");


            yield return WaitForTask(
                nextNormalTask);


            AssertTaskCompletedSuccessfully(
                nextNormalTask,
                "Second persistent conversation run");


            IReadOnlyList<ConversationMessage> nextNormalRequest =
                provider.RequestSnapshots[2];


            Assert.AreEqual(
                3,
                nextNormalRequest.Count);


            Assert.AreEqual(
                "Keep this conversation.",
                nextNormalRequest[0].Content);


            Assert.AreEqual(
                "Response 1",
                nextNormalRequest[1].Content);


            Assert.AreEqual(
                "Continue the persistent conversation.",
                nextNormalRequest[2].Content);
        }


        //--------------------------------------------------
        // Successful Structured Result
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator IsolatedRun_ReturnsSuccessfulStructuredResult()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    new SuccessfulProvider(
                        "Isolated response."));


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


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Successful isolated run");


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


            Assert.IsNotNull(
                result);


            Assert.IsTrue(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


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


            Assert.AreEqual(
                0,
                result.ToolCallsExecuted);


            Assert.IsFalse(
                string.IsNullOrWhiteSpace(
                    result.SessionId));


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


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


        //--------------------------------------------------
        // Failure Preserves Persistent History
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator FailedIsolatedRun_PreservesExistingHistory()
        {
            SequenceProvider provider =
                new SequenceProvider(
                    SequenceProviderMode.SuccessThenFailure);


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


            Task<string> normalTask =
                runtime.ChatAsync(
                    "Persistent successful prompt.");


            yield return WaitForTask(
                normalTask);


            AssertTaskCompletedSuccessfully(
                normalTask,
                "Persistent successful run");


            Assert.AreEqual(
                2,
                runtime.MessageCount);


            Task<AgentRunResult> isolatedTask =
                runtime.RunIsolatedAsync(
                    "This isolated run should fail.");


            yield return WaitForTask(
                isolatedTask);


            AssertTaskCompletedSuccessfully(
                isolatedTask,
                "Failed isolated structured run");


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


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


            StringAssert.Contains(
                "Simulated isolated failure",
                result.ErrorMessage);


            Assert.AreEqual(
                2,
                runtime.MessageCount,
                "A failed isolated run must not affect persistent history.");
        }


        //--------------------------------------------------
        // Cancellation Preserves Persistent History
        //--------------------------------------------------

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


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


            Task<string> normalTask =
                runtime.ChatAsync(
                    "Persistent successful prompt.");


            yield return WaitForTask(
                normalTask);


            AssertTaskCompletedSuccessfully(
                normalTask,
                "Persistent successful run");


            Assert.AreEqual(
                2,
                runtime.MessageCount);


            provider.BlockRequests =
                true;


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<AgentRunResult> isolatedTask =
                runtime.RunIsolatedAsync(
                    "Cancel this isolated run.",
                    cancellation.Token);


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


            Assert.IsTrue(
                runtime.IsRunning);


            cancellation.Cancel();


            yield return WaitForTask(
                isolatedTask);


            AssertTaskCompletedSuccessfully(
                isolatedTask,
                "Cancelled isolated structured run");


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


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsTrue(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


            Assert.AreEqual(
                2,
                runtime.MessageCount,
                "A cancelled isolated run must preserve persistent history.");


            Assert.IsFalse(
                runtime.IsRunning);
        }


        //--------------------------------------------------
        // Timeout Preserves Persistent History
        //--------------------------------------------------

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


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


            Task<string> normalTask =
                runtime.ChatAsync(
                    "Persistent successful prompt.");


            yield return WaitForTask(
                normalTask);


            AssertTaskCompletedSuccessfully(
                normalTask,
                "Persistent successful run");


            Assert.AreEqual(
                2,
                runtime.MessageCount);


            provider.BlockRequests =
                true;


            runtime.ProviderTimeout =
                TimeSpan.FromMilliseconds(
                    100);


            Task<AgentRunResult> isolatedTask =
                runtime.RunIsolatedAsync(
                    "Allow this isolated run to time out.");


            yield return WaitForTask(
                isolatedTask);


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


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


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsTrue(
                result.TimedOut);


            Assert.AreEqual(
                2,
                runtime.MessageCount,
                "A timed-out isolated run must preserve persistent history.");


            Assert.IsFalse(
                runtime.IsRunning);
        }


        //--------------------------------------------------
        // Concurrency Protection
        //--------------------------------------------------

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


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


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<string> activeTask =
                runtime.ChatAsync(
                    "Active persistent run.",
                    cancellation.Token);


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


            Task<AgentRunResult> isolatedTask =
                runtime.RunIsolatedAsync(
                    "This isolated run should be rejected.");


            yield return WaitForTask(
                isolatedTask);


            AssertTaskCompletedSuccessfully(
                isolatedTask,
                "Rejected isolated structured run");


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


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


            Assert.AreEqual(
                "An agent run is already in progress.",
                result.ErrorMessage);


            Assert.AreEqual(
                1,
                runtime.MessageCount,
                "The rejected isolated run must not modify conversation state.");


            cancellation.Cancel();


            yield return WaitForTask(
                activeTask);


            Assert.IsTrue(
                activeTask.IsCanceled);


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


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

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


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


            runtime.RegisterCapability(
                new EchoCapability());


            Task<AgentRunResult> task =
                runtime.RunIsolatedAsync(
                    "Use the echo capability.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Isolated capability run");


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


            Assert.IsTrue(
                result.Succeeded);


            Assert.AreEqual(
                "Capability completed.",
                result.Response);


            Assert.AreEqual(
                1,
                result.ToolCallsExecuted);


            Assert.AreEqual(
                0,
                runtime.MessageCount,
                "Isolated tool-call messages must not enter persistent history.");


            Assert.AreEqual(
                2,
                provider.RequestSnapshots.Count);


            IReadOnlyList<ConversationMessage> secondRequest =
                provider.RequestSnapshots[1];


            Assert.AreEqual(
                3,
                secondRequest.Count,
                "The temporary conversation should contain user, assistant tool call, and tool result.");


            Assert.AreEqual(
                "user",
                secondRequest[0].Role);


            Assert.AreEqual(
                "assistant",
                secondRequest[1].Role);


            Assert.AreEqual(
                "tool",
                secondRequest[2].Role);


            Assert.AreEqual(
                "echo",
                secondRequest[2].Name);


            Assert.AreEqual(
                "hello",
                secondRequest[2].Content);
        }


        //--------------------------------------------------
        // Observer Lifecycle
        //--------------------------------------------------

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


            TestObserver observer =
                new TestObserver();


            runtime.AddAgentObserver(
                observer);


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


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Observed isolated run");


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


            Assert.IsTrue(
                result.Succeeded);


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


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


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


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


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


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

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


            runtime.UseProvider(
                provider);


            return runtime;
        }


        //--------------------------------------------------
        // Test Capability
        //--------------------------------------------------

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


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


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

        private sealed class SuccessfulProvider :
            IProvider
        {
            private readonly string _response;


            public SuccessfulProvider(
                string response)
            {
                _response =
                    response;
            }


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


        private sealed class RecordingProvider :
            IProvider
        {
            private int _requestCount;


            public string Name =>
                "RecordingProvider";


            public List<List<ConversationMessage>> RequestSnapshots
            {
                get;
            } =
                new List<List<ConversationMessage>>();


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


                _requestCount++;


                RequestSnapshots.Add(
                    CloneMessages(
                        messages));


                return Task.FromResult(
                    new ChatCompletion
                    {
                        Content =
                            "Response " +
                            _requestCount
                    });
            }
        }


        private enum SequenceProviderMode
        {
            SuccessThenFailure
        }


        private sealed class SequenceProvider :
            IProvider
        {
            private readonly SequenceProviderMode _mode;

            private int _requestCount;


            public SequenceProvider(
                SequenceProviderMode mode)
            {
                _mode =
                    mode;
            }


            public string Name =>
                "SequenceProvider";


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


                _requestCount++;


                if (
                    _mode == SequenceProviderMode.SuccessThenFailure &&
                    _requestCount == 1)
                {
                    return Task.FromResult(
                        new ChatCompletion
                        {
                            Content =
                                "Persistent response."
                        });
                }


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


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


            public bool BlockRequests
            {
                get;
                set;
            }


            public bool BlockingRequestStarted
            {
                get;
                private set;
            }


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


                if (!BlockRequests)
                {
                    return new ChatCompletion
                    {
                        Content =
                            "Persistent response."
                    };
                }


                BlockingRequestStarted =
                    true;


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


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


        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 CapabilityRequestProvider :
            IProvider
        {
            private bool _capabilityRequested;


            public string Name =>
                "CapabilityRequestProvider";


            public List<List<ConversationMessage>> RequestSnapshots
            {
                get;
            } =
                new List<List<ConversationMessage>>();


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


                RequestSnapshots.Add(
                    CloneMessages(
                        messages));


                if (!_capabilityRequested)
                {
                    _capabilityRequested =
                        true;


                    ChatCompletion completion =
                        new ChatCompletion();


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

                            Name =
                                "echo",

                            Arguments =
                                "{\"Text\":\"hello\"}"
                        });


                    return Task.FromResult(
                        completion);
                }


                return Task.FromResult(
                    new ChatCompletion
                    {
                        Content =
                            "Capability completed."
                    });
            }
        }


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

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


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


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


                return false;
            }
        }


        //--------------------------------------------------
        // Message Snapshot
        //--------------------------------------------------

        private static List<ConversationMessage> CloneMessages(
            IReadOnlyList<ConversationMessage> messages)
        {
            List<ConversationMessage> snapshot =
                new List<ConversationMessage>();


            foreach (ConversationMessage message in messages)
            {
                snapshot.Add(
                    new ConversationMessage
                    {
                        Role =
                            message.Role,

                        Content =
                            message.Content,

                        Name =
                            message.Name,

                        ToolCallId =
                            message.ToolCallId,

                        ToolCalls =
                            message.ToolCalls == null
                                ? null
                                : new List<ToolCall>(
                                    message.ToolCalls)
                    });
            }


            return snapshot;
        }


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