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

using AIRuntime.Runtime.Agents;
using AIRuntime.Runtime.Agents.BuiltIn;
using AIRuntime.Runtime.Agents.Runtime;
using AIRuntime.Runtime.Capabilities;
using AIRuntime.Runtime.Models;
using AIRuntime.Runtime.Providers;
using AIRuntime.Runtime.Providers.Models;
using AIRuntime.Runtime.Services.Agent;

using NUnit.Framework;
using UnityEngine.TestTools;


namespace AIRuntime.Tests
{
    public sealed class AgentRuntimeWiringTests
    {
        private AgentRegistry _registry;

        private AgentSessionManager _sessions;

        private AIRuntime.Runtime.AIRuntime _runtime;



        [SetUp]
        public void Setup()
        {
            _registry =
                new AgentRegistry();


            _sessions =
                new AgentSessionManager(
                    _registry);


            _runtime =
                new AIRuntime.Runtime.AIRuntime(
                    TestPath());
        }



        [UnityTest]
        public IEnumerator Executor_WiresAgentDefinitionIntoRuntime()
        {
            Task<AgentRunResult> task =
                RunWireTest();


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


            Assert.IsFalse(
                task.IsFaulted);


            AgentRunResult result =
                task.Result;


            Assert.AreEqual(
                AgentRunStatus.Succeeded,
                result.Status);
        }

        private async Task<AgentRunResult> RunWireTest()
        {
            AgentDefinition agent =
                ArchitectAgentDefinition.Create();


            _registry.Register(
                agent);


            RecordingProvider provider =
                new RecordingProvider();


            _runtime.UseProvider(
                provider);


            AgentExecutor executor =
                new AgentExecutor(
                    _registry,
                    _sessions,
                    _runtime);


            AgentRunResult result =
                await executor.RunAsync(
                    agent.Id,
                    "Hello");


            Assert.AreEqual(
                agent.SystemPrompt,
                provider.LastSystemPrompt);


            return result;
        }

        [UnityTest]
        public IEnumerator Executor_CreatesIndependentSessions()
        {
            Task test =
                RunSessionTest();


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


            Assert.IsFalse(
                test.IsFaulted);
        }



        private async Task RunSessionTest()
        {
            AgentDefinition agent =
                ArchitectAgentDefinition.Create();


            _registry.Register(
                agent);


            _runtime.UseProvider(
                new RecordingProvider());


            AgentExecutor executor =
                new AgentExecutor(
                    _registry,
                    _sessions,
                    _runtime);


            AgentRunResult first =
                await executor.RunAsync(
                    agent.Id,
                    "First");


            AgentRunResult second =
                await executor.RunAsync(
                    agent.Id,
                    "Second");


            Assert.AreNotEqual(
                first.SessionId,
                second.SessionId);
        }



        [Test]
        public void RuntimeContext_UsesAgentCapabilities()
        {
            AgentDefinition agent =
                ArchitectAgentDefinition.Create();


            _registry.Register(
                agent);


            AgentSession session =
                _sessions.CreateSession(
                    agent.Id);


            AgentExecutionContext execution =
                new AgentExecutionContext(
                    session);


            AgentRuntimeContext context =
                new AgentRuntimeContext(
                    execution,
                    new List<CapabilityDescriptor>());


            Assert.AreEqual(
                agent.Id,
                context.AgentId);


            Assert.AreEqual(
                session.SessionId,
                context.SessionId);
        }



        private static string TestPath()
        {
            return System.IO.Path.Combine(
                UnityEngine.Application.temporaryCachePath,
                "AIRuntimeRuntimeWiringTests");
        }



        //--------------------------------------------------
        // Fake Provider
        //--------------------------------------------------

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


            public string LastSystemPrompt;



            public Task<ChatCompletion> ChatAsync(
                IReadOnlyList<ConversationMessage> messages,
                IReadOnlyList<ProviderTool> tools,
                CancellationToken cancellationToken = default)
            {
                foreach(
                    ConversationMessage message
                    in messages)
                {
                    if(message.Role == "system")
                    {
                        LastSystemPrompt =
                            message.Content;
                    }
                }


                return Task.FromResult(
                    new ChatCompletion
                    {
                        Content =
                            "OK"
                    });
            }
        }
    }
}