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

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


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


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


            Directory.CreateDirectory(
                _workspace);
        }


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


        //--------------------------------------------------
        // Slow Provider Times Out
        //--------------------------------------------------

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


            SlowProvider provider =
                new SlowProvider();


            runtime.ProviderTimeout =
                TimeSpan.FromMilliseconds(
                    100);


            runtime.UseProvider(
                provider);


            Task<string> task =
                runtime.ChatAsync(
                    "This provider should time out.");


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


            Assert.IsTrue(
                task.IsFaulted,
                "ChatAsync should fault with a TimeoutException.");


            Assert.IsFalse(
                task.IsCanceled,
                "A provider timeout should not appear as manual cancellation.");


            Exception exception =
                UnwrapException(
                    task.Exception);


            Assert.IsInstanceOf<TimeoutException>(
                exception);


            StringAssert.Contains(
                "did not respond within",
                exception.Message);


            Assert.IsTrue(
                provider.ReceivedCancellation,
                "The slow provider should receive cancellation from the timeout token.");
        }


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

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


            TestAgentObserver observer =
                new TestAgentObserver();


            runtime.AddAgentObserver(
                observer);


            runtime.ProviderTimeout =
                TimeSpan.FromMilliseconds(
                    100);


            runtime.UseProvider(
                new SlowProvider());


            Task<string> task =
                runtime.ChatAsync(
                    "Emit a timeout event.");


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


            Assert.IsTrue(
                task.IsFaulted);


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


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


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


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


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


            AgentEvent timeoutEvent =
                observer.Find(
                    AgentEventType.RunTimedOut);


            Assert.IsNotNull(
                timeoutEvent);


            Assert.AreEqual(
                runtime.LastRunContext.SessionId,
                timeoutEvent.SessionId);


            StringAssert.Contains(
                "did not respond within",
                timeoutEvent.Message);
        }


        //--------------------------------------------------
        // Structured Timeout Result
        //--------------------------------------------------

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


            runtime.ProviderTimeout =
                TimeSpan.FromMilliseconds(
                    100);


            runtime.UseProvider(
                new SlowProvider());


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Return a timeout result.");


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


            AssertTaskCompletedSuccessfully(
                task,
                "Structured timeout run");


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


            Assert.IsNotNull(
                result);


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsTrue(
                result.TimedOut);


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


            Assert.AreEqual(
                0,
                result.ToolCallsExecuted);


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


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


        //--------------------------------------------------
        // Fast Provider Completes Normally
        //--------------------------------------------------

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


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(
                    1);


            runtime.UseProvider(
                new FastProvider());


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


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


            AssertTaskCompletedSuccessfully(
                task,
                "Fast provider run");


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


            Assert.IsTrue(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


            Assert.AreEqual(
                "Completed before timeout.",
                result.Response);


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


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

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


            SlowProvider provider =
                new SlowProvider();


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(
                    10);


            runtime.UseProvider(
                provider);


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


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


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


            cancellation.Cancel();


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


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


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


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsTrue(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


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


        //--------------------------------------------------
        // Timeout Validation
        //--------------------------------------------------

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


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


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


            Assert.Throws<ArgumentOutOfRangeException>(
                () =>
                {
                    runtime.ProviderTimeout =
                        TimeSpan.FromSeconds(
                            -1);
                });
        }


        //--------------------------------------------------
        // Test Observer
        //--------------------------------------------------

        private sealed class TestAgentObserver :
            IAgentObserver
        {
            public List<AgentEvent> Events
            {
                get;
            } =
                new List<AgentEvent>();


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


            public bool Contains(
                AgentEventType eventType)
            {
                return Find(eventType) != null;
            }


            public AgentEvent Find(
                AgentEventType eventType)
            {
                foreach (AgentEvent agentEvent in Events)
                {
                    if (agentEvent.Type == eventType)
                    {
                        return agentEvent;
                    }
                }


                return null;
            }
        }


        //--------------------------------------------------
        // Slow Provider
        //--------------------------------------------------

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


            public bool HasStarted
            {
                get;
                private set;
            }


            public bool ReceivedCancellation
            {
                get;
                private set;
            }


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


                try
                {
                    await Task.Delay(
                        Timeout.Infinite,
                        cancellationToken);
                }
                catch (OperationCanceledException)
                {
                    ReceivedCancellation =
                        true;


                    throw;
                }


                throw new InvalidOperationException(
                    "The slow provider should never complete normally.");
            }
        }


        //--------------------------------------------------
        // Fast Provider
        //--------------------------------------------------

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


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


                return Task.FromResult(
                    new ChatCompletion
                    {
                        Content =
                            "Completed before timeout."
                    });
            }
        }


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