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

using AIRuntime.Runtime.Capabilities;
using AIRuntime.Runtime.Filesystem;
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 AgentRunContextTests
    {
        private string _workspace;


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


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


            Directory.CreateDirectory(
                _workspace);


            File.WriteAllText(
                Path.Combine(
                    _workspace,
                    "test.txt"),
                "Context lifetime test.");
        }


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


        //--------------------------------------------------
        // Fresh Context Per Run
        //--------------------------------------------------

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


            runtime.UseProvider(
                new ImmediateResponseProvider());


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


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


            AssertTaskCompletedSuccessfully(
                firstTask,
                "First agent run");


            AgentRunContext firstContext =
                runtime.LastRunContext;


            Assert.IsNotNull(
                firstContext);


            Assert.IsTrue(
                firstContext.CompletedAt.HasValue);


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


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


            AssertTaskCompletedSuccessfully(
                secondTask,
                "Second agent run");


            AgentRunContext secondContext =
                runtime.LastRunContext;


            Assert.IsNotNull(
                secondContext);


            Assert.IsTrue(
                secondContext.CompletedAt.HasValue);


            Assert.AreNotSame(
                firstContext,
                secondContext);


            Assert.AreNotEqual(
                firstContext.SessionId,
                secondContext.SessionId);


            Assert.AreEqual(
                0,
                firstContext.ToolCallsExecuted);


            Assert.AreEqual(
                0,
                secondContext.ToolCallsExecuted);
        }


        //--------------------------------------------------
        // Tool Count Resets Per Run
        //--------------------------------------------------

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


            runtime.RegisterCapability(
                new ReadFileCapability());


            runtime.UseProvider(
                new RepeatableReadFileProvider());


            Task<string> firstTask =
                runtime.ChatAsync(
                    "Read the file once.");


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


            AssertTaskCompletedSuccessfully(
                firstTask,
                "First tool run");


            AgentRunContext firstContext =
                runtime.LastRunContext;


            Assert.AreEqual(
                1,
                firstContext.ToolCallsExecuted);


            Task<string> secondTask =
                runtime.ChatAsync(
                    "Read the file again.");


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


            AssertTaskCompletedSuccessfully(
                secondTask,
                "Second tool run");


            AgentRunContext secondContext =
                runtime.LastRunContext;


            Assert.AreNotEqual(
                firstContext.SessionId,
                secondContext.SessionId);


            Assert.AreEqual(
                1,
                firstContext.ToolCallsExecuted);


            Assert.AreEqual(
                1,
                secondContext.ToolCallsExecuted);
        }


        //--------------------------------------------------
        // Events Use Correct Session
        //--------------------------------------------------

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


            TestAgentObserver observer =
                new TestAgentObserver();


            runtime.AddAgentObserver(
                observer);


            runtime.UseProvider(
                new RepeatableReadFileProvider());


            runtime.RegisterCapability(
                new ReadFileCapability());


            Task<string> task =
                runtime.ChatAsync(
                    "Read test.txt");


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


            AssertTaskCompletedSuccessfully(
                task,
                "Lifecycle event run");


            AgentRunContext context =
                runtime.LastRunContext;


            Assert.AreEqual(
                4,
                observer.Events.Count);


            Assert.AreEqual(
                AgentEventType.RunStarted,
                observer.Events[0].Type);


            Assert.AreEqual(
                AgentEventType.ToolStarted,
                observer.Events[1].Type);


            Assert.AreEqual(
                AgentEventType.ToolCompleted,
                observer.Events[2].Type);


            Assert.AreEqual(
                AgentEventType.RunCompleted,
                observer.Events[3].Type);


            foreach (AgentEvent agentEvent in observer.Events)
            {
                Assert.AreEqual(
                    context.SessionId,
                    agentEvent.SessionId);
            }
        }


        //--------------------------------------------------
        // Test Helpers
        //--------------------------------------------------

        private static void AssertTaskCompletedSuccessfully(
            Task task,
            string operation)
        {
            Assert.IsNotNull(
                task);


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


            if (task.IsFaulted)
            {
                System.Exception exception =
                    task.Exception?
                        .Flatten()
                        .InnerException;


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


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


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


        //--------------------------------------------------
        // Provider Returning Immediate Final Answers
        //--------------------------------------------------

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


            public Task<ChatCompletion> ChatAsync(
                IReadOnlyList<ConversationMessage> messages,
                IReadOnlyList<ProviderTool> tools,
                CancellationToken cancellationToken)
            {
                ChatCompletion completion =
                    new ChatCompletion
                    {
                        Content =
                            "Completed."
                    };


                return Task.FromResult(
                    completion);
            }
        }


        //--------------------------------------------------
        // Provider Requesting One Tool Per Run
        //--------------------------------------------------

        private sealed class RepeatableReadFileProvider :
            IProvider
        {
            private int _toolCallNumber;


            public string Name =>
                "RepeatableReadFileProvider";


            public Task<ChatCompletion> ChatAsync(
                IReadOnlyList<ConversationMessage> messages,
                IReadOnlyList<ProviderTool> tools,
                CancellationToken cancellationToken)
            {
                ConversationMessage lastMessage =
                    messages[messages.Count - 1];


                if (lastMessage.Role == "tool")
                {
                    return Task.FromResult(
                        new ChatCompletion
                        {
                            Content =
                                lastMessage.Content
                        });
                }


                _toolCallNumber++;


                ChatCompletion completion =
                    new ChatCompletion();


                completion.ToolCalls.Add(
                    new ToolCall
                    {
                        Id =
                            "context_test_call_" +
                            _toolCallNumber,

                        Name =
                            "read_file",

                        Arguments =
                            "{\"Path\":\"test.txt\"}"
                    });


                return Task.FromResult(
                    completion);
            }
        }
    }
}