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

using AIRuntime.Runtime.Filesystem;
using AIRuntime.Runtime.Providers.Ollama;
using AIRuntime.Runtime.Services.Agent;

using NUnit.Framework;

using UnityEngine;
using UnityEngine.TestTools;


namespace AIRuntime.Tests.Integration
{
    [Category("Integration")]
    [Category("Ollama")]
    public sealed class OllamaAgentLoopTests
    {
        private string _workspace;

        private AIRuntime.Runtime.AIRuntime _runtime;


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


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


            Directory.CreateDirectory(
                _workspace);


            File.WriteAllText(
                Path.Combine(
                    _workspace,
                    "test.txt"),
                "Hello from the Ollama agent test file.");


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


            //--------------------------------------------------
            // Local models can take time to load from disk.
            //--------------------------------------------------

            _runtime.ProviderTimeout =
                TimeSpan.FromMinutes(
                    3);


            _runtime.RunTimeout =
                TimeSpan.FromMinutes(
                    5);


            _runtime.RegisterCapability(
                new ReadFileCapability());


            _runtime.UseProvider(
                new OllamaProvider(
                    "qwen2.5:7b-instruct"));
        }


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


        [UnityTest]
        [Timeout(360000)]
        public IEnumerator OllamaAgentLoop_ReadsFileUsingTool()
        {
            Debug.Log(
                "OLLAMA AGENT LOOP TEST STARTED");


            Task<AgentRunResult> task =
                _runtime.RunAsync(
                    "You must call the read_file tool with path test.txt. " +
                    "After receiving the tool result, return only the exact " +
                    "contents of the file without quotes or explanation.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Ollama agent-loop integration test");


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


            Debug.Log(
                "OLLAMA AGENT LOOP STATUS: " +
                result.Status);


            Debug.Log(
                "OLLAMA AGENT LOOP TOOL CALLS: " +
                result.ToolCallsExecuted);


            Debug.Log(
                "OLLAMA AGENT LOOP FINAL ANSWER:\n" +
                result.Response);


            if (!result.Succeeded)
            {
                Assert.Fail(
                    "The Ollama agent run failed.\n" +
                    "Status: " +
                    result.Status +
                    "\nError: " +
                    result.ErrorMessage);
            }


            Assert.AreEqual(
                1,
                result.ToolCallsExecuted,
                "Ollama should execute read_file exactly once.");


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


            StringAssert.Contains(
                "Hello from the Ollama agent test file.",
                result.Response);
        }


        //--------------------------------------------------
        // Helpers
        //--------------------------------------------------

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


        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 +
                " threw an exception:\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;
        }
    }
}