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


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


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


            Directory.CreateDirectory(
                _workspace);
        }


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


        //--------------------------------------------------
        // Configure Before Provider
        //--------------------------------------------------

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


            AIRuntime.Runtime.AIRuntime runtime =
                new AIRuntime.Runtime.AIRuntime(
                    _workspace);


            runtime.SystemPrompt =
                "You are a Unity coding assistant.";


            runtime.UseProvider(
                provider);


            Task<string> task =
                runtime.ChatAsync(
                    "Inspect this project.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "System prompt run");


            Assert.AreEqual(
                "You are a Unity coding assistant.",
                runtime.SystemPrompt);


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


            IReadOnlyList<ConversationMessage> messages =
                provider.RequestSnapshots[0];


            Assert.AreEqual(
                2,
                messages.Count);


            Assert.AreEqual(
                "system",
                messages[0].Role);


            Assert.AreEqual(
                "You are a Unity coding assistant.",
                messages[0].Content);


            Assert.AreEqual(
                "user",
                messages[1].Role);


            Assert.AreEqual(
                "Inspect this project.",
                messages[1].Content);
        }


        //--------------------------------------------------
        // Configure After Provider
        //--------------------------------------------------

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


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


            runtime.SystemPrompt =
                "Always inspect before editing.";


            Task<string> task =
                runtime.ChatAsync(
                    "Review the scripts.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "System prompt run");


            IReadOnlyList<ConversationMessage> messages =
                provider.RequestSnapshots[0];


            Assert.AreEqual(
                2,
                messages.Count);


            Assert.AreEqual(
                "system",
                messages[0].Role);


            Assert.AreEqual(
                "Always inspect before editing.",
                messages[0].Content);
        }


        //--------------------------------------------------
        // Prompt Appears First
        //--------------------------------------------------

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


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


            runtime.SystemPrompt =
                "Follow the project rules.";


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


            yield return WaitForTask(
                firstTask);


            AssertTaskCompletedSuccessfully(
                firstTask,
                "First conversation run");


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


            yield return WaitForTask(
                secondTask);


            AssertTaskCompletedSuccessfully(
                secondTask,
                "Second conversation run");


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


            Assert.AreEqual(
                4,
                messages.Count);


            Assert.AreEqual(
                "system",
                messages[0].Role);


            Assert.AreEqual(
                "Follow the project rules.",
                messages[0].Content);


            Assert.AreEqual(
                "user",
                messages[1].Role);


            Assert.AreEqual(
                "First prompt.",
                messages[1].Content);


            Assert.AreEqual(
                "assistant",
                messages[2].Role);


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


            Assert.AreEqual(
                "user",
                messages[3].Role);


            Assert.AreEqual(
                "Second prompt.",
                messages[3].Content);
        }


        //--------------------------------------------------
        // Persistent Run
        //--------------------------------------------------

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


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


            runtime.SystemPrompt =
                "Persistent system instructions.";


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Run a persistent task.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Persistent structured run");


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


            Assert.IsTrue(
                result.Succeeded);


            IReadOnlyList<ConversationMessage> messages =
                provider.RequestSnapshots[0];


            Assert.AreEqual(
                "system",
                messages[0].Role);


            Assert.AreEqual(
                "Persistent system instructions.",
                messages[0].Content);
        }


        //--------------------------------------------------
        // Isolated Run
        //--------------------------------------------------

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


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


            runtime.SystemPrompt =
                "Isolated runs must follow these rules.";


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


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Isolated structured run");


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


            Assert.IsTrue(
                result.Succeeded);


            IReadOnlyList<ConversationMessage> messages =
                provider.RequestSnapshots[0];


            Assert.AreEqual(
                2,
                messages.Count);


            Assert.AreEqual(
                "system",
                messages[0].Role);


            Assert.AreEqual(
                "Isolated runs must follow these rules.",
                messages[0].Content);


            Assert.AreEqual(
                "user",
                messages[1].Role);


            Assert.AreEqual(
                "Run an isolated task.",
                messages[1].Content);


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


        //--------------------------------------------------
        // Reset Preserves Prompt
        //--------------------------------------------------

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


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


            runtime.SystemPrompt =
                "Keep this system prompt.";


            Task<string> firstTask =
                runtime.ChatAsync(
                    "Store a conversation.");


            yield return WaitForTask(
                firstTask);


            AssertTaskCompletedSuccessfully(
                firstTask,
                "Conversation run");


            Assert.AreEqual(
                2,
                runtime.MessageCount);


            runtime.ResetConversation();


            Assert.AreEqual(
                0,
                runtime.MessageCount);


            Assert.AreEqual(
                "Keep this system prompt.",
                runtime.SystemPrompt);


            Task<string> secondTask =
                runtime.ChatAsync(
                    "Start a clean conversation.");


            yield return WaitForTask(
                secondTask);


            AssertTaskCompletedSuccessfully(
                secondTask,
                "Post-reset conversation run");


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


            Assert.AreEqual(
                2,
                messages.Count);


            Assert.AreEqual(
                "system",
                messages[0].Role);


            Assert.AreEqual(
                "Keep this system prompt.",
                messages[0].Content);


            Assert.AreEqual(
                "Start a clean conversation.",
                messages[1].Content);
        }


        //--------------------------------------------------
        // Clear Prompt
        //--------------------------------------------------

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


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


            runtime.SystemPrompt =
                "Remove these instructions.";


            runtime.ClearSystemPrompt();


            Assert.AreEqual(
                string.Empty,
                runtime.SystemPrompt);


            Task<string> task =
                runtime.ChatAsync(
                    "Run without a system prompt.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Prompt-cleared run");


            IReadOnlyList<ConversationMessage> messages =
                provider.RequestSnapshots[0];


            Assert.AreEqual(
                1,
                messages.Count);


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


            Assert.AreEqual(
                "Run without a system prompt.",
                messages[0].Content);
        }


        //--------------------------------------------------
        // Clear Before Provider
        //--------------------------------------------------

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


            runtime.SystemPrompt =
                "Temporary prompt.";


            Assert.DoesNotThrow(
                runtime.ClearSystemPrompt);


            Assert.AreEqual(
                string.Empty,
                runtime.SystemPrompt);
        }


        //--------------------------------------------------
        // Empty Prompt
        //--------------------------------------------------

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


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


            runtime.SystemPrompt =
                "   ";


            Task<string> task =
                runtime.ChatAsync(
                    "Only send the user message.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Empty-system-prompt run");


            Assert.AreEqual(
                string.Empty,
                runtime.SystemPrompt);


            IReadOnlyList<ConversationMessage> messages =
                provider.RequestSnapshots[0];


            Assert.AreEqual(
                1,
                messages.Count);


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


        //--------------------------------------------------
        // Prompt Is Not Persistent History
        //--------------------------------------------------

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


            runtime.SystemPrompt =
                "Do not count this message.";


            Assert.AreEqual(
                0,
                runtime.MessageCount);


            Task<string> task =
                runtime.ChatAsync(
                    "Count this user message.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Message-count run");


            Assert.AreEqual(
                2,
                runtime.MessageCount,
                "Only the user and assistant messages should be stored.");
        }


        //--------------------------------------------------
        // Provider Switching
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ProviderSwitch_PreservesSystemPrompt()
        {
            RecordingProvider firstProvider =
                new RecordingProvider();


            RecordingProvider secondProvider =
                new RecordingProvider();


            AIRuntime.Runtime.AIRuntime runtime =
                new AIRuntime.Runtime.AIRuntime(
                    _workspace);


            runtime.SystemPrompt =
                "Preserve this prompt between providers.";


            runtime.UseProvider(
                firstProvider);


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


            yield return WaitForTask(
                firstTask);


            AssertTaskCompletedSuccessfully(
                firstTask,
                "First provider run");


            runtime.UseProvider(
                secondProvider);


            Assert.AreEqual(
                "Preserve this prompt between providers.",
                runtime.SystemPrompt);


            Assert.AreEqual(
                0,
                runtime.MessageCount,
                "Changing providers should still reset conversation history.");


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


            yield return WaitForTask(
                secondTask);


            AssertTaskCompletedSuccessfully(
                secondTask,
                "Second provider run");


            IReadOnlyList<ConversationMessage> messages =
                secondProvider.RequestSnapshots[0];


            Assert.AreEqual(
                2,
                messages.Count);


            Assert.AreEqual(
                "system",
                messages[0].Role);


            Assert.AreEqual(
                "Preserve this prompt between providers.",
                messages[0].Content);


            Assert.AreEqual(
                "Second provider request.",
                messages[1].Content);
        }


        //--------------------------------------------------
        // Running-State Protection
        //--------------------------------------------------

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


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


            runtime.SystemPrompt =
                "Original prompt.";


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


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


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


            InvalidOperationException exception =
                Assert.Throws<InvalidOperationException>(
                    () =>
                    {
                        runtime.SystemPrompt =
                            "Changed prompt.";
                    });


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


            Assert.AreEqual(
                "Original prompt.",
                runtime.SystemPrompt);


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            Assert.IsTrue(
                task.IsCanceled);
        }


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


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


            runtime.SystemPrompt =
                "Original prompt.";


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


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


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


            InvalidOperationException exception =
                Assert.Throws<InvalidOperationException>(
                    runtime.ClearSystemPrompt);


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


            Assert.AreEqual(
                "Original prompt.",
                runtime.SystemPrompt);


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            Assert.IsTrue(
                task.IsCanceled);
        }


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

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


            runtime.UseProvider(
                provider);


            return runtime;
        }


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

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


        //--------------------------------------------------
        // 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 complete normally.");
            }
        }


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