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

using AIRuntime.Runtime.Models;
using AIRuntime.Runtime.Providers;
using AIRuntime.Runtime.Providers.Models;
using AIRuntime.Runtime.Providers.Ollama;
using AIRuntime.Runtime.Providers.Testing;

using NUnit.Framework;

using UnityEngine.TestTools;


namespace AIRuntime.Tests
{
    public sealed class ProviderContractTests
    {
        //--------------------------------------------------
        // Provider Identity
        //--------------------------------------------------

        [Test]
        public void Provider_HasName()
        {
            IProvider provider =
                new FakeProvider();


            Assert.IsFalse(
                string.IsNullOrWhiteSpace(
                    provider.Name));
        }


        //--------------------------------------------------
        // Real Ollama Provider
        //--------------------------------------------------

        [UnityTest]
        [Category("Integration")]
        public IEnumerator OllamaProvider_ReturnsResponse()
        {
            OllamaProvider provider =
                new OllamaProvider(
                    "qwen2.5:7b-instruct");


            List<ConversationMessage> messages =
                new List<ConversationMessage>
                {
                    new ConversationMessage
                    {
                        Role =
                            "user",

                        Content =
                            "Say hello."
                    }
                };


            Task<ChatCompletion> task =
                provider.ChatAsync(
                    messages,
                    new List<ProviderTool>());


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


            AssertTaskCompletedSuccessfully(
                task,
                "Ollama provider request");


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


            Assert.IsNotNull(
                result);


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


        //--------------------------------------------------
        // Provider Cancellation
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator OllamaProvider_RejectsCancelledRequest()
        {
            OllamaProvider provider =
                new OllamaProvider(
                    "qwen2.5:7b-instruct");


            List<ConversationMessage> messages =
                new List<ConversationMessage>
                {
                    new ConversationMessage
                    {
                        Role =
                            "user",

                        Content =
                            "This request should never reach Ollama."
                    }
                };


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            cancellation.Cancel();


            Task<ChatCompletion> task =
                provider.ChatAsync(
                    messages,
                    new List<ProviderTool>(),
                    cancellation.Token);


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


            Assert.IsTrue(
                task.IsCanceled,
                "The provider task should be cancelled.");


            Assert.IsFalse(
                task.IsFaulted,
                "Cancellation should not be reported as a provider failure.");
        }

//--------------------------------------------------
// AIRuntime Provider Configuration
//--------------------------------------------------

        [UnityTest]
        public IEnumerator UseProvider_ConfiguresRuntime()
        {
            string workspace =
                System.IO.Path.Combine(
                    UnityEngine.Application.temporaryCachePath,
                    "AIRuntimeUseProviderTest");


            System.IO.Directory.CreateDirectory(
                workspace);


            try
            {
                AIRuntime.Runtime.AIRuntime runtime =
                    new AIRuntime.Runtime.AIRuntime(
                        workspace);


                runtime.UseProvider(
                    new FinalResponseProvider());


                Task<string> task =
                    runtime.ChatAsync(
                        "Test provider configuration.");


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


                AssertTaskCompletedSuccessfully(
                    task,
                    "Configured runtime request");


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


                Assert.AreEqual(
                    "Provider configured successfully.",
                    result);
            }
            finally
            {
                if (System.IO.Directory.Exists(workspace))
                {
                    System.IO.Directory.Delete(
                        workspace,
                        true);
                }
            }
        }
        //--------------------------------------------------
        // Provider Can Receive Tools
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator Provider_AcceptsToolDefinitions()
        {
            IProvider provider =
                new FakeProvider();


            List<ProviderTool> tools =
                new List<ProviderTool>
                {
                    new ProviderTool
                    {
                        Type =
                            "function",

                        Function =
                            new ProviderToolFunction
                            {
                                Name =
                                    "read_file",

                                Description =
                                    "Reads a file.",

                                Parameters =
                                    new object()
                            }
                    }
                };


            List<ConversationMessage> messages =
                new List<ConversationMessage>
                {
                    new ConversationMessage
                    {
                        Role =
                            "user",

                        Content =
                            "Read test.txt"
                    }
                };


            Task<ChatCompletion> task =
                provider.ChatAsync(
                    messages,
                    tools);


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


            AssertTaskCompletedSuccessfully(
                task,
                "Fake provider tool-definition request");


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


            Assert.IsNotNull(
                result);
        }


        //--------------------------------------------------
        // Provider Can Request Tools
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator Provider_CanReturnToolCalls()
        {
            IProvider provider =
                new FakeProvider();


            Task<ChatCompletion> task =
                provider.ChatAsync(
                    new List<ConversationMessage>(),
                    new List<ProviderTool>());


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


            AssertTaskCompletedSuccessfully(
                task,
                "Fake provider tool-call request");


            ChatCompletion completion =
                task.GetAwaiter()
                    .GetResult();


            Assert.IsNotNull(
                completion);


            Assert.IsTrue(
                completion.HasToolCalls);


            Assert.IsNotEmpty(
                completion.ToolCalls);


            Assert.AreEqual(
                "read_file",
                completion.ToolCalls[0].Name);


            Assert.IsFalse(
                string.IsNullOrWhiteSpace(
                    completion.ToolCalls[0].Arguments));
        }


        //--------------------------------------------------
        // Provider Can Return Final Answer
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator Provider_CanReturnFinalResponse()
        {
            IProvider provider =
                new FakeProvider();


            // First call consumes the configured tool response.

            Task<ChatCompletion> firstTask =
                provider.ChatAsync(
                    new List<ConversationMessage>(),
                    new List<ProviderTool>());


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


            AssertTaskCompletedSuccessfully(
                firstTask,
                "Fake provider initial tool-call request");


            ChatCompletion firstCompletion =
                firstTask.GetAwaiter()
                    .GetResult();


            Assert.IsNotNull(
                firstCompletion);


            Assert.IsTrue(
                firstCompletion.HasToolCalls);


            // Second call returns the final response.

            Task<ChatCompletion> secondTask =
                provider.ChatAsync(
                    new List<ConversationMessage>(),
                    new List<ProviderTool>());


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


            AssertTaskCompletedSuccessfully(
                secondTask,
                "Fake provider final-response request");


            ChatCompletion finalCompletion =
                secondTask.GetAwaiter()
                    .GetResult();


            Assert.IsNotNull(
                finalCompletion);


            Assert.IsFalse(
                finalCompletion.HasToolCalls);


            Assert.IsFalse(
                string.IsNullOrWhiteSpace(
                    finalCompletion.Content));
        }


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

        private static void AssertTaskCompletedSuccessfully(
            Task task,
            string operation)
        {
            Assert.IsNotNull(
                task,
                operation + " returned a null task.");


            if (task.IsCanceled)
            {
                Assert.Fail(
                    operation + " was cancelled.");
            }


            if (!task.IsFaulted)
                return;


            Exception exception =
                UnwrapException(
                    task.Exception);


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

//--------------------------------------------------
// Provider Returning a Final Response
//--------------------------------------------------

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


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


                return Task.FromResult(
                    new ChatCompletion
                    {
                        Content =
                            "Provider configured successfully."
                    });
            }
        }
        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;
        }
    }
}