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

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

using RuntimeCore = AIRuntime.Runtime;

using NUnit.Framework;
using UnityEngine.TestTools;


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

        private AgentSessionManager _sessionManager;

        private RuntimeCore.AIRuntime _runtime;


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


            _sessionManager =
                new AgentSessionManager(
                    _registry);


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



        [Test]
        public void ExecutionContext_CreatesFromSession()
        {
            AgentDefinition agent =
                CreateAgent();


            _registry.Register(
                agent);


            AgentSession session =
                _sessionManager.CreateSession(
                    "TestAgent");


            AgentExecutionContext context =
                new AgentExecutionContext(
                    session);


            Assert.AreSame(
                session,
                context.Session);


            Assert.AreSame(
                agent,
                context.Definition);
        }



        [Test]
        public void ExecutionContext_UsesDefinitionIdentity()
        {
            AgentDefinition agent =
                CreateAgent();


            _registry.Register(
                agent);


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


            AgentExecutionContext context =
                new AgentExecutionContext(
                    session);


            Assert.AreEqual(
                "TestAgent",
                context.AgentId);


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



        [Test]
        public void Executor_UnknownAgentThrows()
        {
            AgentExecutor executor =
                CreateExecutor();


            AggregateException exception =
                Assert.Throws<AggregateException>(
                    () =>
                    {
                        executor.RunAsync(
                                "MissingAgent",
                                "Hello")
                            .Wait();
                    });


            Assert.IsInstanceOf<InvalidOperationException>(
                exception.InnerException);
        }


        [UnityTest]
        public IEnumerator Executor_UsesAgentSystemPrompt()
        {
            AgentDefinition agent =
                CreateAgent();

            _registry.Register(
                agent);


            RecordingProvider provider =
                new RecordingProvider();


            _runtime.UseProvider(
                provider);


            AgentExecutor executor =
                CreateExecutor();


            Task<AgentRunResult> task =
                executor.RunAsync(
                    "TestAgent",
                    "Hello");


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


            Assert.AreEqual(
                "Agent system prompt",
                provider.LastSystemPrompt);
        }



        [UnityTest]
        public IEnumerator Executor_RestoresPreviousPrompt()
        {
            AgentDefinition agent =
                CreateAgent();


            _registry.Register(
                agent);


            _runtime.SystemPrompt =
                "Original prompt";


            RecordingProvider provider =
                new RecordingProvider();


            _runtime.UseProvider(
                provider);


            AgentExecutor executor =
                CreateExecutor();


            Task<AgentRunResult> task =
                executor.RunAsync(
                    "TestAgent",
                    "Hello");


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


            Assert.AreEqual(
                "Original prompt",
                _runtime.SystemPrompt);
        }



        private AgentExecutor CreateExecutor()
        {
            return new AgentExecutor(
                _registry,
                _sessionManager,
                _runtime);
        }



        private static AgentDefinition CreateAgent()
        {
            return new AgentDefinition(
                "TestAgent",
                "Test Agent",
                "Agent execution test",
                "Agent system prompt");
        }



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



        //--------------------------------------------------
        // 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 =
                            "Test response"
                    });
            }
        }
    }
}