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

using Newtonsoft.Json.Linq;

using AIRuntime.Runtime.Capabilities;
using AIRuntime.Runtime.Capabilities.Attributes;
using AIRuntime.Runtime.Capabilities.Permissions;
using AIRuntime.Runtime.Filesystem;
using AIRuntime.Runtime.Providers.Models;
using AIRuntime.Runtime.Tools;

using NUnit.Framework;

using UnityEngine;
using UnityEngine.TestTools;


namespace AIRuntime.Tests
{
    public sealed class CapabilityTests
    {
        private AIRuntime.Runtime.AIRuntime _runtime;

        private string _workspace;


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


            Directory.CreateDirectory(
                _workspace);


            File.WriteAllText(
                Path.Combine(
                    _workspace,
                    "test.txt"),
                "AIRuntime test file contents.");


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


            _runtime.RegisterCapability(
                new ReadFileCapability());


            _runtime.RegisterCapability(
                new WriteFileCapability());


            _runtime.RegisterCapability(
                new AppendFileCapability());


            //--------------------------------------------------
            // Capability Behaviour Test Configuration
            //--------------------------------------------------
            //
            // These tests validate filesystem capability behaviour,
            // validation and workspace security.
            //
            // Permission enforcement itself is covered separately by
            // CapabilityPermissionTests and CapabilityApprovalTests.
            //--------------------------------------------------

            _runtime.AllowPermission(
                CapabilityPermission.Write);
        }


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


        //--------------------------------------------------
        // Read
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ReadFile_ReturnsContents()
        {
            Task<string> task =
                _runtime.ExecuteCapabilityAsync(
                    "read_file",
                    @"{
                        ""path"": ""test.txt""
                    }");


            yield return AwaitTask(
                task);


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


            Assert.IsTrue(
                result.Contains(
                    "AIRuntime"));
        }


        [UnityTest]
        public IEnumerator ReadFile_RecoversUniqueEncodedModelPath()
        {
            string docsDirectory =
                Path.Combine(
                    _workspace,
                    "Assets",
                    "Runtime",
                    "Docs");

            Directory.CreateDirectory(
                docsDirectory);

            File.WriteAllText(
                Path.Combine(
                    docsDirectory,
                    "DEPENDENCY_RULES.md"),
                "ConversationEngine is the only orchestrator.");


            Task<string> task =
                _runtime.ExecuteCapabilityAsync(
                    "read_file",
                    @"{
                        ""path"": ""Assets_Runtime_Docs_DEPENDENCY_RULES_md""
                    }");


            yield return AwaitTask(
                task);


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


            StringAssert.Contains(
                "only orchestrator",
                result);
        }

        [UnityTest]
        public IEnumerator ReadFile_RecoversEncodedDirectoryWithExactFileName()
        {
            string docsDirectory =
                Path.Combine(
                    _workspace,
                    "Assets",
                    "Runtime",
                    "Docs");

            Directory.CreateDirectory(
                docsDirectory);

            File.WriteAllText(
                Path.Combine(
                    docsDirectory,
                    "DEPENDENCY_RULES.md"),
                "Providers must never depend on Modules.");


            Task<string> task =
                _runtime.ExecuteCapabilityAsync(
                    "read_file",
                    @"{
                        ""path"": ""Assets_Runtime_Docs/DEPENDENCY_RULES.md""
                    }");


            yield return AwaitTask(
                task);


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


            StringAssert.Contains(
                "must never depend on Modules",
                result);
        }


        [UnityTest]
        public IEnumerator ReadFile_RecoversDuplicatedFinalExtension()
        {
            string materialsDirectory =
                Path.Combine(
                    _workspace,
                    "Assets",
                    "Materials");

            Directory.CreateDirectory(
                materialsDirectory);

            File.WriteAllText(
                Path.Combine(
                    materialsDirectory,
                    "GoalMat.mat"),
                "m_Name: GoalMat");


            Task<string> task =
                _runtime.ExecuteCapabilityAsync(
                    "read_file",
                    @"{
                        ""path"": ""Assets/Materials/GoalMat.mat.mat""
                    }");


            yield return AwaitTask(
                task);


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

            StringAssert.Contains(
                "m_Name: GoalMat",
                result);
        }


        [UnityTest]
        public IEnumerator ReadFile_DoesNotRecoverUniqueFileNameFromWrongDirectory()
        {
            string scriptsDirectory =
                Path.Combine(
                    _workspace,
                    "Assets",
                    "JumpingBallGame",
                    "Scripts");

            Directory.CreateDirectory(
                scriptsDirectory);

            File.WriteAllText(
                Path.Combine(
                    scriptsDirectory,
                    "SideScrollCameraFollow.cs"),
                "public sealed class SideScrollCameraFollow {}");


            LogAssert.Expect(
                LogType.Error,
                "[Capability] 'read_file' failed: File not found: Assets/Runtime/Scripts/SideScrollCameraFollow.cs");

            yield return ExpectExceptionAsync<FileNotFoundException>(
                () =>
                    _runtime.ExecuteCapabilityAsync(
                        "read_file",
                        @"{
                            ""path"": ""Assets/Runtime/Scripts/SideScrollCameraFollow.cs""
                        }"));
        }


        [UnityTest]
        public IEnumerator ReadFile_RequireExactPathDoesNotRecoverWrongDirectory()
        {
            string scriptsDirectory =
                Path.Combine(
                    _workspace,
                    "Assets",
                    "JumpingBallGame",
                    "Scripts");

            Directory.CreateDirectory(
                scriptsDirectory);

            File.WriteAllText(
                Path.Combine(
                    scriptsDirectory,
                    "SideScrollCameraFollow.cs"),
                "public sealed class SideScrollCameraFollow {}");

            LogAssert.Expect(
                LogType.Error,
                "[Capability] 'read_file' failed: File not found: Assets/Runtime/Scripts/SideScrollCameraFollow.cs");

            yield return ExpectExceptionAsync<FileNotFoundException>(
                () =>
                    _runtime.ExecuteCapabilityAsync(
                        "read_file",
                        @"{
                            ""path"": ""Assets/Runtime/Scripts/SideScrollCameraFollow.cs"",
                            ""requireExactPath"": true
                        }"));
        }


        [UnityTest]
        public IEnumerator ReadFile_DoesNotRecoverAmbiguousFileName()
        {
            foreach (string directoryName in new[]
                     {
                         "First",
                         "Second"
                     })
            {
                string directory =
                    Path.Combine(
                        _workspace,
                        "Assets",
                        directoryName);

                Directory.CreateDirectory(
                    directory);

                File.WriteAllText(
                    Path.Combine(
                        directory,
                        "Shared.cs"),
                    directoryName);
            }


            LogAssert.Expect(
                LogType.Error,
                "[Capability] 'read_file' failed: File not found: Assets/Wrong/Shared.cs");


            yield return ExpectExceptionAsync<FileNotFoundException>(
                () =>
                    _runtime.ExecuteCapabilityAsync(
                        "read_file",
                        @"{
                            ""path"": ""Assets/Wrong/Shared.cs""
                        }"));
        }


        //--------------------------------------------------
        // Write
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator WriteFile_CreatesFile()
        {
            Task<string> task =
                _runtime.ExecuteCapabilityAsync(
                    "write_file",
                    @"{
                        ""path"": ""generated/test.txt"",
                        ""content"": ""Hello AIRuntime""
                    }");


            yield return AwaitTask(
                task);


            string path =
                Path.Combine(
                    _workspace,
                    "generated",
                    "test.txt");


            Assert.IsTrue(
                File.Exists(
                    path));


            Assert.AreEqual(
                "Hello AIRuntime",
                File.ReadAllText(
                    path));
        }


        [UnityTest]
        public IEnumerator WriteFile_OverwritesExistingFile()
        {
            string path =
                Path.Combine(
                    _workspace,
                    "overwrite.txt");


            File.WriteAllText(
                path,
                "old");


            Task<string> task =
                _runtime.ExecuteCapabilityAsync(
                    "write_file",
                    @"{
                        ""path"": ""overwrite.txt"",
                        ""content"": ""new""
                    }");


            yield return AwaitTask(
                task);


            Assert.AreEqual(
                "new",
                File.ReadAllText(
                    path));
        }


        //--------------------------------------------------
        // Append
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator AppendFile_Appends()
        {
            string path =
                Path.Combine(
                    _workspace,
                    "append.txt");


            File.WriteAllText(
                path,
                "A");


            Task<string> task =
                _runtime.ExecuteCapabilityAsync(
                    "append_file",
                    @"{
                        ""path"": ""append.txt"",
                        ""content"": ""B""
                    }");


            yield return AwaitTask(
                task);


            Assert.AreEqual(
                "A\nB",
                File.ReadAllText(
                    path));
        }


        [UnityTest]
        public IEnumerator AppendFile_DoesNotDuplicateNewLineBoundary()
        {
            string path =
                Path.Combine(
                    _workspace,
                    "append-boundary.txt");

            File.WriteAllText(
                path,
                "A\n");


            yield return AwaitTask(
                _runtime.ExecuteCapabilityAsync(
                    "append_file",
                    @"{
                        ""path"": ""append-boundary.txt"",
                        ""content"": ""B""
                    }"));

            yield return AwaitTask(
                _runtime.ExecuteCapabilityAsync(
                    "append_file",
                    @"{
                        ""path"": ""append-boundary.txt"",
                        ""content"": ""\nC""
                    }"));


            Assert.AreEqual(
                "A\nB\nC",
                File.ReadAllText(
                    path));
        }


        [Test]
        public void AppendFile_AdvertisesExactMutationContract()
        {
            AppendFileCapability capability =
                new AppendFileCapability();

            CapabilitySchemaExporter exporter =
                new CapabilitySchemaExporter();

            JObject schema =
                exporter.Export(
                    capability);


            StringAssert.Contains(
                "must call this tool",
                capability.Description);

            StringAssert.Contains(
                "Preserve all slashes",
                schema["properties"]?["path"]?["description"]?
                    .Value<string>());

            StringAssert.Contains(
                "newline is added automatically",
                schema["properties"]?["content"]?["description"]?
                    .Value<string>());
        }


        [UnityTest]
        public IEnumerator AppendFile_CreatesMissingFile()
        {
            Task<string> task =
                _runtime.ExecuteCapabilityAsync(
                    "append_file",
                    @"{
                        ""path"": ""new.txt"",
                        ""content"": ""created""
                    }");


            yield return AwaitTask(
                task);


            Assert.AreEqual(
                "created",
                File.ReadAllText(
                    Path.Combine(
                        _workspace,
                        "new.txt")));
        }


        [UnityTest]
        public IEnumerator AppendFile_CreatesDirectories()
        {
            Task<string> task =
                _runtime.ExecuteCapabilityAsync(
                    "append_file",
                    @"{
                        ""path"": ""logs/runtime/log.txt"",
                        ""content"": ""entry""
                    }");


            yield return AwaitTask(
                task);


            Assert.IsTrue(
                File.Exists(
                    Path.Combine(
                        _workspace,
                        "logs",
                        "runtime",
                        "log.txt")));
        }


        //--------------------------------------------------
        // Security
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ReadFile_PreventsEscape()
        {
            LogAssert.Expect(
                LogType.Error,
                "[Capability] 'read_file' failed: Path '../../../Windows/System32/config.sys' escapes the workspace.");


            yield return ExpectExceptionAsync<Exception>(
                () =>
                    _runtime.ExecuteCapabilityAsync(
                        "read_file",
                        @"{
                            ""path"": ""../../../Windows/System32/config.sys""
                        }"));
        }


        [UnityTest]
        public IEnumerator WriteFile_PreventsEscape()
        {
            LogAssert.Expect(
                LogType.Error,
                "[Capability] 'write_file' failed: Path '../../../escape.txt' escapes the workspace.");


            yield return ExpectExceptionAsync<Exception>(
                () =>
                    _runtime.ExecuteCapabilityAsync(
                        "write_file",
                        @"{
                            ""path"": ""../../../escape.txt"",
                            ""content"": ""blocked""
                        }"));
        }


        [UnityTest]
        public IEnumerator AppendFile_PreventsEscape()
        {
            LogAssert.Expect(
                LogType.Error,
                "[Capability] 'append_file' failed: Path '../../../escape.txt' escapes the workspace.");


            yield return ExpectExceptionAsync<Exception>(
                () =>
                    _runtime.ExecuteCapabilityAsync(
                        "append_file",
                        @"{
                            ""path"": ""../../../escape.txt"",
                            ""content"": ""blocked""
                        }"));
        }


        //--------------------------------------------------
        // Validation
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator UnknownCapabilityThrows()
        {
            yield return ExpectExceptionAsync<InvalidOperationException>(
                () =>
                    _runtime.ExecuteCapabilityAsync(
                        "missing",
                        "{}"));
        }


        [UnityTest]
        public IEnumerator MissingArgumentsThrows()
        {
            LogAssert.Expect(
                LogType.Warning,
                "[Capability] 'read_file' failed: Path is required.");


            yield return ExpectExceptionAsync<ValidationException>(
                () =>
                    _runtime.ExecuteCapabilityAsync(
                        "read_file",
                        "{}"));
        }


        [UnityTest]
        public IEnumerator InvalidJsonThrows()
        {
            yield return ExpectExceptionAsync<Exception>(
                () =>
                    _runtime.ExecuteCapabilityAsync(
                        "read_file",
                        "{"));
        }


        //--------------------------------------------------
        // Capability Discovery
        //--------------------------------------------------

        [Test]
        public void DiscoversCapabilities()
        {
            var capabilities =
                _runtime.GetCapabilities();


            Assert.AreEqual(
                3,
                capabilities.Count);


            Assert.IsTrue(
                capabilities.Any(
                    capability =>
                        capability.Name ==
                        "read_file"));
        }


        [Test]
        public void CapabilityHasSchema()
        {
            ReadFileCapability capability =
                new ReadFileCapability();


            CapabilitySchemaExporter exporter =
                new CapabilitySchemaExporter();


            JObject schema =
                exporter.Export(
                    capability);


            Assert.IsNotNull(
                schema);


            Assert.IsTrue(
                schema["properties"]?["path"] != null);

            Assert.AreEqual(
                "Exact workspace-relative path of the existing text file whose contents must be read. Preserve '/' directory separators exactly; never replace separators with underscores.",
                schema["properties"]?["path"]?["description"]?
                    .Value<string>());

            StringAssert.Contains(
                "must call this tool",
                capability.Description);


            CollectionAssert.Contains(
                schema["required"]
                    .Values<string>()
                    .ToArray(),
                "path");


            Assert.IsFalse(
                schema["properties"]?["Path"] != null);


            CollectionAssert.DoesNotContain(
                schema["required"]
                    .Values<string>()
                    .ToArray(),
                "Path");
        }


        [Test]
        public void ProviderExport_PreservesGeneratedSchemaDetails()
        {
            _runtime.RegisterCapability(
                new RichSchemaCapability());


            ProviderTool tool =
                _runtime.ExportProviderTools()
                    .First(
                        item =>
                            item.Function.Name ==
                            "rich_schema");


            JObject schema =
                JObject.FromObject(
                    tool.Function.Parameters);


            Assert.AreEqual(
                "integer",
                schema["properties"]?["count"]?["type"]?
                    .Value<string>());


            Assert.AreEqual(
                "Number of attempts.",
                schema["properties"]?["count"]?["description"]?
                    .Value<string>());


            Assert.AreEqual(
                3,
                schema["properties"]?["count"]?["default"]?
                    .Value<int>());


            CollectionAssert.AreEqual(
                new[] { "Fast", "Safe" },
                schema["properties"]?["mode"]?["enum"]?
                    .Values<string>()
                    .ToArray());


            Assert.AreEqual(
                "string",
                schema["properties"]?["tags"]?["items"]?["type"]?
                    .Value<string>());


            Assert.AreEqual(
                "string",
                schema["properties"]?["options"]?["properties"]?["label"]?["type"]?
                    .Value<string>());
        }


        //--------------------------------------------------
        // Provider Tool Export
        //--------------------------------------------------

        [Test]
        public void ExportsProviderTools()
        {
            var tools =
                _runtime.ExportProviderTools();


            Assert.AreEqual(
                3,
                tools.Count);


            Assert.IsTrue(
                tools.Any(
                    tool =>
                        tool.Function.Name ==
                        "read_file"));
        }


        [Test]
        public void ProviderToolsContainSchema()
        {
            var tool =
                _runtime.ExportProviderTools()
                    .First(
                        item =>
                            item.Function.Name ==
                            "read_file");


            Assert.IsNotNull(
                tool.Function.Parameters);


            Assert.IsNotNull(
                tool.Function.Description);
        }


        [Test]
        public void ProviderToolsAreFunctions()
        {
            foreach (
                ProviderTool tool
                in _runtime.ExportProviderTools())
            {
                Assert.AreEqual(
                    "function",
                    tool.Type);
            }
        }


        //--------------------------------------------------
        // OpenAI Schema
        //--------------------------------------------------

        [Test]
        public void ExportToolSchemas_ReturnsOpenAIFormat()
        {
            JArray tools =
                _runtime.ExportToolSchemas();


            Assert.AreEqual(
                3,
                tools.Count);


            Assert.AreEqual(
                "function",
                tools[0]["type"]
                    .Value<string>());
        }


        //--------------------------------------------------
        // Tool Parser
        //--------------------------------------------------

        [Test]
        public void ToolCallParserReadsName()
        {
            ToolCallParser parser =
                new ToolCallParser();


            ToolCallRequest request =
                parser.Parse(
                    @"{
                        ""name"": ""read_file"",
                        ""arguments"": {
                            ""path"": ""test.txt""
                        }
                    }");


            Assert.AreEqual(
                "read_file",
                request.Name);


            Assert.AreEqual(
                "test.txt",
                request.Arguments["path"]
                    .Value<string>());
        }


        //--------------------------------------------------
        // Tool Executor
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ToolExecutorExecutesCapability()
        {
            Task<ToolCallResult> task =
                _runtime.ExecuteToolCallAsync(
                    @"{
                        ""name"": ""read_file"",
                        ""arguments"": {
                            ""path"": ""test.txt""
                        }
                    }");


            yield return AwaitTask(
                task);


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


            Assert.IsTrue(
                result.Success);


            StringAssert.Contains(
                "AIRuntime",
                result.Output);
        }


        //--------------------------------------------------
        // End To End
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator EndToEndWriteAppendRead()
        {
            yield return AwaitTask(
                _runtime.ExecuteCapabilityAsync(
                    "write_file",
                    @"{
                        ""path"": ""e2e/output.txt"",
                        ""content"": ""start""
                    }"));


            yield return AwaitTask(
                _runtime.ExecuteCapabilityAsync(
                    "append_file",
                    @"{
                        ""path"": ""e2e/output.txt"",
                        ""content"": ""finish""
                    }"));


            Task<string> readTask =
                _runtime.ExecuteCapabilityAsync(
                    "read_file",
                    @"{
                        ""path"": ""e2e/output.txt""
                    }");


            yield return AwaitTask(
                readTask);


            Assert.AreEqual(
                "start\nfinish",
                readTask.GetAwaiter()
                    .GetResult());
        }


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

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


            if (task.IsFaulted)
            {
                throw UnwrapException(
                    task.Exception);
            }


            if (task.IsCanceled)
            {
                throw new OperationCanceledException(
                    "The test task was unexpectedly cancelled.");
            }
        }


        private static IEnumerator ExpectExceptionAsync<T>(
            Func<Task> operation)
            where T : Exception
        {
            Task task =
                null;


            Exception exception =
                null;


            try
            {
                task =
                    operation();
            }
            catch (Exception caughtException)
            {
                exception =
                    caughtException;
            }


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


                if (task.IsFaulted)
                {
                    exception =
                        UnwrapException(
                            task.Exception);
                }
                else if (task.IsCanceled)
                {
                    exception =
                        new OperationCanceledException(
                            "The operation was cancelled.");
                }
            }


            Assert.IsNotNull(
                exception);


            Assert.IsInstanceOf<T>(
                exception);
        }


        private static Exception UnwrapException(
            Exception exception)
        {
            if (exception is AggregateException aggregate)
            {
                aggregate =
                    aggregate.Flatten();


                if (aggregate.InnerExceptions.Count == 1)
                {
                    return aggregate.InnerExceptions[0];
                }
            }


            return exception;
        }


        private enum RichSchemaMode
        {
            Fast,
            Safe
        }


        private sealed class RichSchemaOptions
        {
            public string Label
            {
                get;
                set;
            }
        }


        private sealed class RichSchemaArguments
        {
            [AIRuntime.Runtime.Capabilities.Attributes.Description(
                "Number of attempts.")]
            [DefaultValue(3)]
            public int Count
            {
                get;
                set;
            }


            public RichSchemaMode Mode
            {
                get;
                set;
            }


            public string[] Tags
            {
                get;
                set;
            }


            public RichSchemaOptions Options
            {
                get;
                set;
            }
        }


        private sealed class RichSchemaCapability
            : CapabilityBase<RichSchemaArguments>
        {
            public override string Name =>
                "rich_schema";


            public override string Description =>
                "Exercises rich schema export.";


            public override CapabilityPermission Permission =>
                CapabilityPermission.ReadOnly;


            protected override Task<object> ExecuteAsync(
                RichSchemaArguments arguments,
                CapabilityContext context)
            {
                return Task.FromResult<object>(
                    null);
            }
        }
    }
}
