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 CapabilityApprovalTests
    {
        private const string WriteArgumentsJson =
            "{\"value\":\"approved content\"}";


        private string _workspace;


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


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


            Directory.CreateDirectory(
                _workspace);
        }


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


        //--------------------------------------------------
        // Handler Configuration
        //--------------------------------------------------

        [Test]
        public void ApprovalHandler_IsNotConfiguredByDefault()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            Assert.IsFalse(
                runtime.HasApprovalHandler);
        }


        [Test]
        public void SetApprovalHandler_ConfiguresHandler()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            RecordingApprovalHandler handler =
                RecordingApprovalHandler.CreateApproved();


            runtime.SetApprovalHandler(
                handler);


            Assert.IsTrue(
                runtime.HasApprovalHandler);
        }


        [Test]
        public void ClearApprovalHandler_RemovesHandler()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            runtime.SetApprovalHandler(
                RecordingApprovalHandler.CreateApproved());


            Assert.IsTrue(
                runtime.HasApprovalHandler);


            runtime.ClearApprovalHandler();


            Assert.IsFalse(
                runtime.HasApprovalHandler);
        }


        [UnityTest]
        public IEnumerator MultiEffectCapability_RequiresEveryDeclaredPermission()
        {
            CountingMultiEffectCapability capability = new CountingMultiEffectCapability();
            AIRuntime.Runtime.AIRuntime runtime = CreateRuntime();
            runtime.RegisterCapability(capability);
            runtime.AllowPermission(CapabilityPermission.Write);

            Task<string> denied = runtime.ExecuteCapabilityAsync(capability.Name, "{}");
            yield return WaitForTask(denied);
            Assert.IsTrue(denied.IsFaulted);
            Assert.AreEqual(0, capability.ExecutionCount);

            runtime.AllowPermission(CapabilityPermission.Network);
            Task<string> allowed = runtime.ExecuteCapabilityAsync(capability.Name, "{}");
            yield return WaitForTask(allowed);
            Assert.IsFalse(allowed.IsFaulted);
            Assert.AreEqual(1, capability.ExecutionCount);
        }


        //--------------------------------------------------
        // Missing Handler
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ApprovalRequired_WithoutHandler_DoesNotExecute()
        {
            CountingWriteCapability capability =
                new CountingWriteCapability();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            runtime.RegisterCapability(
                capability);


            Task<string> task =
                runtime.ExecuteCapabilityAsync(
                    capability.Name,
                    WriteArgumentsJson);


            yield return WaitForTask(
                task);


            Assert.IsTrue(
                task.IsFaulted);


            CapabilityPermissionException exception =
                AssertPermissionException(
                    task.Exception);


            Assert.AreEqual(
                CapabilityPermissionDecision.RequireApproval,
                exception.Decision);


            StringAssert.Contains(
                "no approval handler is configured",
                exception.Message);


            Assert.AreEqual(
                0,
                capability.ExecutionCount);
        }


        //--------------------------------------------------
        // Approval
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ApprovalHandler_CanApproveCapability()
        {
            CountingWriteCapability capability =
                new CountingWriteCapability();


            RecordingApprovalHandler handler =
                RecordingApprovalHandler.CreateApproved(
                    "Approved by test.");


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            runtime.RegisterCapability(
                capability);


            runtime.SetApprovalHandler(
                handler);


            Task<string> task =
                runtime.ExecuteCapabilityAsync(
                    capability.Name,
                    WriteArgumentsJson);


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Approved capability");


            Assert.AreEqual(
                "Write completed.",
                task.Result);


            Assert.AreEqual(
                1,
                capability.ExecutionCount);


            Assert.AreEqual(
                1,
                handler.RequestCount);
        }


        //--------------------------------------------------
        // Rejection
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ApprovalHandler_CanRejectCapability()
        {
            CountingWriteCapability capability =
                new CountingWriteCapability();


            RecordingApprovalHandler handler =
                RecordingApprovalHandler.CreateRejected(
                    "The operation was rejected by the user.");


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            runtime.RegisterCapability(
                capability);


            runtime.SetApprovalHandler(
                handler);


            Task<string> task =
                runtime.ExecuteCapabilityAsync(
                    capability.Name,
                    WriteArgumentsJson);


            yield return WaitForTask(
                task);


            Assert.IsTrue(
                task.IsFaulted);


            CapabilityPermissionException exception =
                AssertPermissionException(
                    task.Exception);


            Assert.AreEqual(
                CapabilityPermissionDecision.RequireApproval,
                exception.Decision);


            StringAssert.Contains(
                "was rejected during approval",
                exception.Message);


            StringAssert.Contains(
                "The operation was rejected by the user.",
                exception.Message);


            Assert.AreEqual(
                0,
                capability.ExecutionCount);


            Assert.AreEqual(
                1,
                handler.RequestCount);
        }


        //--------------------------------------------------
        // Null Handler Result
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator NullApprovalResult_IsRejectedSafely()
        {
            CountingWriteCapability capability =
                new CountingWriteCapability();


            RecordingApprovalHandler handler =
                RecordingApprovalHandler.CreateNullResult();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            runtime.RegisterCapability(
                capability);


            runtime.SetApprovalHandler(
                handler);


            Task<string> task =
                runtime.ExecuteCapabilityAsync(
                    capability.Name,
                    WriteArgumentsJson);


            yield return WaitForTask(
                task);


            Assert.IsTrue(
                task.IsFaulted);


            CapabilityPermissionException exception =
                AssertPermissionException(
                    task.Exception);


            StringAssert.Contains(
                "returned no result",
                exception.Message);


            Assert.AreEqual(
                0,
                capability.ExecutionCount);
        }


        //--------------------------------------------------
        // Request Data
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ApprovalRequest_ContainsCapabilityMetadata()
        {
            CountingWriteCapability capability =
                new CountingWriteCapability();


            RecordingApprovalHandler handler =
                RecordingApprovalHandler.CreateApproved();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            runtime.RegisterCapability(
                capability);


            runtime.SetApprovalHandler(
                handler);


            Task<string> task =
                runtime.ExecuteCapabilityAsync(
                    capability.Name,
                    WriteArgumentsJson);


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Approval metadata test");


            CapabilityApprovalRequest request =
                handler.LastRequest;


            Assert.IsNotNull(
                request);


            Assert.AreEqual(
                capability.Name,
                request.CapabilityName);


            Assert.AreEqual(
                CapabilityPermission.Write,
                request.Permission);


            Assert.AreEqual(
                WriteArgumentsJson,
                request.Arguments);


            Assert.AreEqual(
                string.Empty,
                request.SessionId,
                "Direct capability execution should not have an agent session ID.");
        }


        //--------------------------------------------------
        // Agent Session
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator AgentApprovalRequest_ContainsRunSessionId()
        {
            CountingWriteCapability capability =
                new CountingWriteCapability();


            RecordingApprovalHandler handler =
                RecordingApprovalHandler.CreateApproved();


            ToolThenResponseProvider provider =
                new ToolThenResponseProvider(
                    capability.Name,
                    WriteArgumentsJson);


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


            runtime.RegisterCapability(
                capability);


            runtime.SetApprovalHandler(
                handler);


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


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Agent approval session test");


            AgentRunResult result =
                task.Result;


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


            Assert.AreEqual(
                1,
                result.ToolCallsExecuted);


            Assert.AreEqual(
                1,
                capability.ExecutionCount);


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


            Assert.IsNotNull(
                handler.LastRequest);


            Assert.AreEqual(
                result.SessionId,
                handler.LastRequest.SessionId);


            Assert.AreEqual(
                WriteArgumentsJson,
                handler.LastRequest.Arguments);
        }


        //--------------------------------------------------
        // Agent Rejection
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator RejectedAgentCapability_ReturnsFailedStatus()
        {
            CountingWriteCapability capability =
                new CountingWriteCapability();


            RecordingApprovalHandler handler =
                RecordingApprovalHandler.CreateRejected(
                    "User selected Reject.");


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    new SingleToolCallProvider(
                        capability.Name,
                        WriteArgumentsJson));


            runtime.RegisterCapability(
                capability);


            runtime.SetApprovalHandler(
                handler);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Attempt the rejected write capability.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Rejected agent capability");


            AgentRunResult result =
                task.Result;


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


            Assert.IsTrue(
                result.Failed);


            StringAssert.Contains(
                "User selected Reject.",
                result.ErrorMessage);


            Assert.AreEqual(
                0,
                result.ToolCallsExecuted);


            Assert.AreEqual(
                0,
                capability.ExecutionCount);


            Assert.AreEqual(
                0,
                runtime.MessageCount,
                "The failed agent run should roll back its conversation.");
        }


        //--------------------------------------------------
        // Cancellation
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ApprovalCancellation_CancelsDirectExecution()
        {
            CountingWriteCapability capability =
                new CountingWriteCapability();


            BlockingApprovalHandler handler =
                new BlockingApprovalHandler();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            runtime.RegisterCapability(
                capability);


            runtime.SetApprovalHandler(
                handler);


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<string> task =
                runtime.ExecuteCapabilityAsync(
                    capability.Name,
                    WriteArgumentsJson,
                    cancellation.Token);


            yield return WaitForApprovalStart(
                handler);


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            Assert.IsTrue(
                task.IsCanceled ||
                UnwrapException(task.Exception) is OperationCanceledException);


            Assert.AreEqual(
                0,
                capability.ExecutionCount);
        }


        [UnityTest]
        public IEnumerator ApprovalCancellation_ReturnsCancelledAgentStatus()
        {
            CountingWriteCapability capability =
                new CountingWriteCapability();


            BlockingApprovalHandler handler =
                new BlockingApprovalHandler();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    new SingleToolCallProvider(
                        capability.Name,
                        WriteArgumentsJson));


            runtime.RegisterCapability(
                capability);


            runtime.SetApprovalHandler(
                handler);


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


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


            yield return WaitForApprovalStart(
                handler);


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Cancelled approval agent run");


            AgentRunResult result =
                task.Result;


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


            Assert.AreEqual(
                0,
                result.ToolCallsExecuted);


            Assert.AreEqual(
                0,
                capability.ExecutionCount);


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


        //--------------------------------------------------
        // Handler Lock
        //--------------------------------------------------

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


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


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


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


            yield return WaitForProviderStart(
                provider);


            InvalidOperationException exception =
                Assert.Throws<InvalidOperationException>(
                    () =>
                    {
                        runtime.SetApprovalHandler(
                            RecordingApprovalHandler.CreateApproved());
                    });


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


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Approval-handler lock test");


            Assert.AreEqual(
                AgentRunStatus.Cancelled,
                task.Result.Status);
        }


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

        private AIRuntime.Runtime.AIRuntime CreateRuntime()
        {
            return CreateRuntime(
                new SuccessfulProvider());
        }


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


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(10);


            runtime.RunTimeout =
                TimeSpan.FromSeconds(10);


            runtime.UseProvider(
                provider);


            return runtime;
        }


        //--------------------------------------------------
        // Capability
        //--------------------------------------------------

        private sealed class CountingWriteCapability :
            CapabilityBase<WriteArguments>
        {
            public int ExecutionCount
            {
                get;
                private set;
            }


            public override string Name =>
                "approval_write_test";


            public override string Description =>
                "Writes test content after approval.";


            public override CapabilityPermission Permission =>
                CapabilityPermission.Write;


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


                ExecutionCount++;


                return Task.FromResult<object>(
                    "Write completed.");
            }
        }


        private sealed class CountingMultiEffectCapability :
            CapabilityBase<WriteArguments>, ICapabilityAuthorizationSemantics
        {
            public int ExecutionCount { get; private set; }
            public override string Name => "multi_effect_test";
            public override string Description => "Tests multi-effect authorization.";
            public override CapabilityPermission Permission => CapabilityPermission.Write;
            public IReadOnlyList<CapabilityPermission> RequiredPermissions =>
                new[] { CapabilityPermission.Write, CapabilityPermission.Network };
            public IReadOnlyList<string> EffectScopes => new[] { "test:workspace", "test:network" };
            protected override Task<object> ExecuteAsync(WriteArguments arguments, CapabilityContext context)
            {
                ExecutionCount++;
                return Task.FromResult<object>("completed");
            }
        }


        private sealed class WriteArguments
        {
            public string Value
            {
                get;
                set;
            }
        }


        //--------------------------------------------------
        // Approval Handlers
        //--------------------------------------------------

        private sealed class RecordingApprovalHandler :
            ICapabilityApprovalHandler
        {
            private readonly CapabilityApprovalResult _result;


            private RecordingApprovalHandler(
                CapabilityApprovalResult result)
            {
                _result =
                    result;
            }


            public int RequestCount
            {
                get;
                private set;
            }


            public CapabilityApprovalRequest LastRequest
            {
                get;
                private set;
            }


            public static RecordingApprovalHandler CreateApproved(
                string reason = null)
            {
                return new RecordingApprovalHandler(
                    CapabilityApprovalResult.Approve(
                        reason));
            }


            public static RecordingApprovalHandler CreateRejected(
                string reason = null)
            {
                return new RecordingApprovalHandler(
                    CapabilityApprovalResult.Reject(
                        reason));
            }


            public static RecordingApprovalHandler CreateNullResult()
            {
                return new RecordingApprovalHandler(
                    null);
            }


            public Task<CapabilityApprovalResult> RequestApprovalAsync(
                CapabilityApprovalRequest request,
                CancellationToken cancellationToken)
            {
                cancellationToken
                    .ThrowIfCancellationRequested();


                RequestCount++;


                LastRequest =
                    request;


                return Task.FromResult(
                    _result);
            }
        }


        private sealed class BlockingApprovalHandler :
            ICapabilityApprovalHandler
        {
            public bool HasStarted
            {
                get;
                private set;
            }


            public async Task<CapabilityApprovalResult> RequestApprovalAsync(
                CapabilityApprovalRequest request,
                CancellationToken cancellationToken)
            {
                HasStarted =
                    true;


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


                return CapabilityApprovalResult.Approve();
            }
        }


        //--------------------------------------------------
        // 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 SingleToolCallProvider :
            IProvider
        {
            private readonly string _capabilityName;

            private readonly string _arguments;


            public SingleToolCallProvider(
                string capabilityName,
                string arguments)
            {
                _capabilityName =
                    capabilityName;


                _arguments =
                    arguments;
            }


            public string Name =>
                "SingleToolCallProvider";


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


                return Task.FromResult(
                    CreateToolCompletion(
                        _capabilityName,
                        _arguments,
                        "approval_call"));
            }
        }


        private sealed class ToolThenResponseProvider :
            IProvider
        {
            private readonly string _capabilityName;

            private readonly string _arguments;

            private bool _toolRequested;


            public ToolThenResponseProvider(
                string capabilityName,
                string arguments)
            {
                _capabilityName =
                    capabilityName;


                _arguments =
                    arguments;
            }


            public string Name =>
                "ToolThenResponseProvider";


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


                if (!_toolRequested)
                {
                    _toolRequested =
                        true;


                    return Task.FromResult(
                        CreateToolCompletion(
                            _capabilityName,
                            _arguments,
                            "approved_agent_call"));
                }


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


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


        private static ChatCompletion CreateToolCompletion(
            string capabilityName,
            string arguments,
            string callId)
        {
            ChatCompletion completion =
                new ChatCompletion();


            completion.ToolCalls.Add(
                new ToolCall
                {
                    Id =
                        callId,

                    Name =
                        capabilityName,

                    Arguments =
                        arguments
                });


            return completion;
        }


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

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


        private static IEnumerator WaitForApprovalStart(
            BlockingApprovalHandler handler)
        {
            while (!handler.HasStarted)
            {
                yield return null;
            }
        }


        private static IEnumerator WaitForProviderStart(
            BlockingProvider provider)
        {
            while (!provider.HasStarted)
            {
                yield return null;
            }
        }


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

        private static CapabilityPermissionException AssertPermissionException(
            Exception exception)
        {
            Exception unwrapped =
                UnwrapException(
                    exception);


            Assert.IsInstanceOf<CapabilityPermissionException>(
                unwrapped);


            return (CapabilityPermissionException)unwrapped;
        }


        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 == null)
                return null;


            if (!(exception is AggregateException aggregate))
                return exception;


            AggregateException flattened =
                aggregate.Flatten();


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


            return flattened;
        }
    }
}
