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

using AIRuntime.Runtime.Filesystem;
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 AgentRunResultTests
    {
        private string _workspace;


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


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


            Directory.CreateDirectory(
                _workspace);


            File.WriteAllText(
                Path.Combine(
                    _workspace,
                    "test.txt"),
                "Structured result test file.");
        }


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


        //--------------------------------------------------
        // Successful Run
        //--------------------------------------------------

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


            runtime.UseProvider(
                new SuccessfulProvider());


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Return a successful response.");


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


            AssertTaskCompletedSuccessfully(
                task,
                "Successful structured run");


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


            Assert.IsNotNull(
                result);


            Assert.IsTrue(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


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


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


            Assert.AreEqual(
                0,
                result.ToolCallsExecuted);


            Assert.GreaterOrEqual(
                result.Duration,
                System.TimeSpan.Zero);


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


        //--------------------------------------------------
        // Tool Count
        //--------------------------------------------------

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


            runtime.RegisterCapability(
                new ReadFileCapability());


            runtime.UseProvider(
                new ReadFileProvider());


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Read test.txt");


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


            AssertTaskCompletedSuccessfully(
                task,
                "Structured tool run");


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


            Assert.IsTrue(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


            Assert.AreEqual(
                1,
                result.ToolCallsExecuted);


            Assert.AreEqual(
                "Structured result test file.",
                result.Response);


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


        //--------------------------------------------------
        // Cancelled Run
        //--------------------------------------------------

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


            BlockingProvider provider =
                new BlockingProvider();


            runtime.UseProvider(
                provider);


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Wait until cancelled.",
                    cancellation.Token);


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


            cancellation.Cancel();


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


            AssertTaskCompletedSuccessfully(
                task,
                "Cancelled structured run");


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


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsTrue(
                result.Cancelled);


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


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


            Assert.AreEqual(
                0,
                result.ToolCallsExecuted);


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


        //--------------------------------------------------
        // Failed Run
        //--------------------------------------------------

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


            runtime.UseProvider(
                new FailingProvider());


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Trigger a failure.");


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


            AssertTaskCompletedSuccessfully(
                task,
                "Failed structured run");


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


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


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


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


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


        //--------------------------------------------------
        // ChatAsync Compatibility
        //--------------------------------------------------

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


            runtime.UseProvider(
                new SuccessfulProvider());


            Task<string> task =
                runtime.ChatAsync(
                    "Return the normal response.");


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


            AssertTaskCompletedSuccessfully(
                task,
                "ChatAsync compatibility run");


            string response =
                task.GetAwaiter()
                    .GetResult();


            Assert.AreEqual(
                "Structured response.",
                response);
        }


        //--------------------------------------------------
        // Early Failure Context Isolation
        //--------------------------------------------------

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

            runtime.UseProvider(
                new SuccessfulProvider());


            Task<AgentRunResult> acceptedTask =
                runtime.RunAsync(
                    "Complete the accepted run.");

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

            AssertTaskCompletedSuccessfully(
                acceptedTask,
                "Accepted structured run");


            AgentRunResult acceptedResult =
                acceptedTask.GetAwaiter()
                    .GetResult();

            AgentRunContext acceptedContext =
                runtime.LastRunContext;


            Task<AgentRunResult> rejectedTask =
                runtime.RunAsync(
                    "   ");

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

            AssertTaskCompletedSuccessfully(
                rejectedTask,
                "Validation failure");


            AgentRunResult rejectedResult =
                rejectedTask.GetAwaiter()
                    .GetResult();


            Assert.IsTrue(
                rejectedResult.Failed);

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

            Assert.AreNotEqual(
                acceptedResult.SessionId,
                rejectedResult.SessionId);

            Assert.AreEqual(
                0,
                rejectedResult.ToolCallsExecuted);

            Assert.AreSame(
                acceptedContext,
                runtime.LastRunContext);
        }


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

            runtime.UseProvider(
                new SuccessfulProvider());


            Task<AgentRunResult> acceptedTask =
                runtime.RunAsync(
                    "Complete the accepted run.");

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

            AssertTaskCompletedSuccessfully(
                acceptedTask,
                "Accepted structured run");


            AgentRunResult acceptedResult =
                acceptedTask.GetAwaiter()
                    .GetResult();

            AgentRunContext acceptedContext =
                runtime.LastRunContext;


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();

            cancellation.Cancel();


            Task<AgentRunResult> cancelledTask =
                runtime.RunAsync(
                    "Do not start this run.",
                    cancellation.Token);

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

            AssertTaskCompletedSuccessfully(
                cancelledTask,
                "Pre-cancelled structured run");


            AgentRunResult cancelledResult =
                cancelledTask.GetAwaiter()
                    .GetResult();


            Assert.IsTrue(
                cancelledResult.Cancelled);

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

            Assert.AreNotEqual(
                acceptedResult.SessionId,
                cancelledResult.SessionId);

            Assert.AreEqual(
                0,
                cancelledResult.ToolCallsExecuted);

            Assert.AreSame(
                acceptedContext,
                runtime.LastRunContext);
        }


        //--------------------------------------------------
        // Task Validation
        //--------------------------------------------------

        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;


            System.Exception exception =
                task.Exception?
                    .Flatten()
                    .InnerException;


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


        //--------------------------------------------------
        // Successful Provider
        //--------------------------------------------------

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


        //--------------------------------------------------
        // Tool Provider
        //--------------------------------------------------

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


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


                ConversationMessage lastMessage =
                    messages[messages.Count - 1];


                if (lastMessage.Role == "tool")
                {
                    return Task.FromResult(
                        new ChatCompletion
                        {
                            Content =
                                lastMessage.Content
                        });
                }


                ChatCompletion completion =
                    new ChatCompletion();


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

                        Name =
                            "read_file",

                        Arguments =
                            "{\"Path\":\"test.txt\"}"
                    });


                return Task.FromResult(
                    completion);
            }
        }


        //--------------------------------------------------
        // Blocking Provider
        //--------------------------------------------------

        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 System.InvalidOperationException(
                    "Blocking provider completed unexpectedly.");
            }
        }


        //--------------------------------------------------
        // Failing Provider
        //--------------------------------------------------

        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 System.InvalidOperationException(
                        "Expected structured provider failure."));
            }
        }
    }
}
