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


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


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


            Directory.CreateDirectory(
                _workspace);
        }


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


        //--------------------------------------------------
        // Reset Conversation
        //--------------------------------------------------

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


            Task<string> task =
                runtime.ChatAsync(
                    "First prompt.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "First conversation run");


            Assert.AreEqual(
                2,
                runtime.MessageCount,
                "A successful run should store one user and one assistant message.");


            runtime.ResetConversation();


            Assert.AreEqual(
                0,
                runtime.MessageCount,
                "ResetConversation should remove all stored messages.");
        }


        //--------------------------------------------------
        // Successful History
        //--------------------------------------------------

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


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


            Task<string> firstTask =
                runtime.ChatAsync(
                    "First prompt.");


            yield return WaitForTask(
                firstTask);


            AssertTaskCompletedSuccessfully(
                firstTask,
                "First conversation run");


            Assert.AreEqual(
                2,
                runtime.MessageCount);


            Task<string> secondTask =
                runtime.ChatAsync(
                    "Second prompt.");


            yield return WaitForTask(
                secondTask);


            AssertTaskCompletedSuccessfully(
                secondTask,
                "Second conversation run");


            Assert.AreEqual(
                4,
                runtime.MessageCount,
                "Two successful runs should preserve four messages.");


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


            Assert.AreEqual(
                1,
                provider.RequestMessageCounts[0],
                "The first provider request should receive the first user message.");


            Assert.AreEqual(
                3,
                provider.RequestMessageCounts[1],
                "The second provider request should receive previous history plus the new user message.");


            Assert.AreEqual(
                "First prompt.",
                provider.RequestSnapshots[1][0].Content);


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


            Assert.AreEqual(
                "Second prompt.",
                provider.RequestSnapshots[1][2].Content);
        }


        //--------------------------------------------------
        // Cancellation Rollback
        //--------------------------------------------------

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


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


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "This prompt should be rolled back.",
                    cancellation.Token);


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


            Assert.AreEqual(
                1,
                runtime.MessageCount,
                "The user message should exist while the run is active.");


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Cancelled structured run");


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


            Assert.IsTrue(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


            Assert.AreEqual(
                0,
                runtime.MessageCount,
                "A cancelled run should roll back its user message.");
        }


        //--------------------------------------------------
        // Timeout Rollback
        //--------------------------------------------------

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


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


            runtime.ProviderTimeout =
                TimeSpan.FromMilliseconds(
                    100);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "This prompt should time out and roll back.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Timed-out structured run");


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


            Assert.IsTrue(
                result.TimedOut);


            Assert.IsFalse(
                result.Cancelled);


            Assert.AreEqual(
                0,
                runtime.MessageCount,
                "A timed-out run should roll back its user message.");
        }


        //--------------------------------------------------
        // Failure Rollback
        //--------------------------------------------------

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


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "This prompt should fail and roll back.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Failed structured run");


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


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


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


            Assert.AreEqual(
                0,
                runtime.MessageCount,
                "A failed run should roll back its user message.");
        }


        //--------------------------------------------------
        // Preserve Existing History During Rollback
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator FailedRun_PreservesEarlierSuccessfulHistory()
        {
            SequenceProvider provider =
                new SequenceProvider();


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


            Task<string> firstTask =
                runtime.ChatAsync(
                    "Keep this successful exchange.");


            yield return WaitForTask(
                firstTask);


            AssertTaskCompletedSuccessfully(
                firstTask,
                "Successful conversation run");


            Assert.AreEqual(
                2,
                runtime.MessageCount);


            Task<AgentRunResult> failedTask =
                runtime.RunAsync(
                    "Roll back only this failed exchange.");


            yield return WaitForTask(
                failedTask);


            AssertTaskCompletedSuccessfully(
                failedTask,
                "Failed structured run");


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


            Assert.IsFalse(
                result.Succeeded);


            Assert.AreEqual(
                2,
                runtime.MessageCount,
                "Rollback should preserve all messages before the checkpoint.");
        }


        //--------------------------------------------------
        // Provider Switch
        //--------------------------------------------------

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


            Task<string> firstTask =
                runtime.ChatAsync(
                    "Message for the first provider.");


            yield return WaitForTask(
                firstTask);


            AssertTaskCompletedSuccessfully(
                firstTask,
                "First provider run");


            Assert.AreEqual(
                2,
                runtime.MessageCount);


            RecordingProvider secondProvider =
                new RecordingProvider();


            runtime.UseProvider(
                secondProvider);


            Assert.AreEqual(
                0,
                runtime.MessageCount,
                "Changing providers should create a clean conversation.");


            Task<string> secondTask =
                runtime.ChatAsync(
                    "Message for the second provider.");


            yield return WaitForTask(
                secondTask);


            AssertTaskCompletedSuccessfully(
                secondTask,
                "Second provider run");


            Assert.AreEqual(
                1,
                secondProvider.RequestMessageCounts[0],
                "The new provider should receive only its new user prompt.");
        }


        //--------------------------------------------------
        // Reset Without Provider
        //--------------------------------------------------

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


            Assert.DoesNotThrow(
                runtime.ResetConversation);


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


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

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


            runtime.UseProvider(
                provider);


            return runtime;
        }


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


        //--------------------------------------------------
        // Recording Provider
        //--------------------------------------------------

        private sealed class RecordingProvider :
            IProvider
        {
            private int _responseNumber;


            public string Name =>
                "RecordingProvider";


            public List<int> RequestMessageCounts
            {
                get;
            } =
                new List<int>();


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


                _responseNumber++;


                RequestMessageCounts.Add(
                    messages.Count);


                RequestSnapshots.Add(
                    CloneMessages(
                        messages));


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


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


        //--------------------------------------------------
        // 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 InvalidOperationException(
                    "The blocking provider should not finish normally.");
            }
        }


        //--------------------------------------------------
        // 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();


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


        //--------------------------------------------------
        // Success Then Failure Provider
        //--------------------------------------------------

        private sealed class SequenceProvider :
            IProvider
        {
            private int _requestCount;


            public string Name =>
                "SequenceProvider";


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


                _requestCount++;


                if (_requestCount == 1)
                {
                    return Task.FromResult(
                        new ChatCompletion
                        {
                            Content =
                                "Stored successful response."
                        });
                }


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


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