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

using AIRuntime.Runtime;
using AIRuntime.Runtime.Models;
using AIRuntime.Runtime.Providers;
using AIRuntime.Runtime.Providers.Fallback;
using AIRuntime.Runtime.Providers.Health;
using AIRuntime.Runtime.Providers.Models;

using NUnit.Framework;

using UnityEngine;
using UnityEngine.TestTools;


namespace AIRuntime.Tests
{
    public sealed class ProviderHealthFallbackTests
    {
        private string _workspace;


        private AIRuntime.Runtime.AIRuntime _runtime;



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


            System.IO.Directory.CreateDirectory(
                _workspace);


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


            _runtime.ProviderRetryCount =
                1;
        }



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



        //--------------------------------------------------
        // Retry
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator FailedProvider_RetriesBeforeFallback()
        {
            FakeProvider primary =
                new FakeProvider(
                    "Primary",
                    failCount: 1);


            FakeProvider fallback =
                new FakeProvider(
                    "Fallback",
                    failCount: 0);



            _runtime.UseProvider(
                primary);


            _runtime.AddFallbackProvider(
                fallback);



            Task<string> task =
                _runtime.ChatAsync(
                    "hello");


            yield return AwaitTask(
                task);



            Assert.AreEqual(
                2,
                primary.CallCount);


            Assert.AreEqual(
                0,
                fallback.CallCount);
        }



        //--------------------------------------------------
        // Cooldown skip
        //--------------------------------------------------

     [UnityTest]
public IEnumerator UnhealthyProvider_IsSkipped()
{
    FakeProvider primary =
        new FakeProvider(
            "Primary",
            failCount: 99);


    FakeProvider fallback =
        new FakeProvider(
            "Fallback",
            failCount: 0);


    //--------------------------------------------------
    // One initial attempt plus one retry gives the
    // primary two failures during the first run.
    //--------------------------------------------------

    _runtime.ProviderRetryCount =
        1;


    _runtime.ProviderRetryDelay =
        TimeSpan.Zero;


    _runtime.ProviderMaximumRetryDelay =
        TimeSpan.Zero;


    _runtime.ProviderCooldownFailureThreshold =
        2;


    _runtime.ProviderBaseCooldown =
        TimeSpan.FromSeconds(
            30);


    _runtime.ProviderMaximumCooldown =
        TimeSpan.FromSeconds(
            30);


    _runtime.UseProvider(
        primary);


    _runtime.AddFallbackProvider(
        fallback);


    //--------------------------------------------------
    // First run:
    // primary fails twice and enters cooldown;
    // fallback succeeds.
    //--------------------------------------------------

    Task<string> first =
        _runtime.ChatAsync(
            "first");


    yield return AwaitTask(
        first);


    Assert.AreEqual(
        "Fallback response",
        first.GetAwaiter()
            .GetResult());


    Assert.AreEqual(
        2,
        primary.CallCount,
        "The primary should receive one attempt and one retry.");


    Assert.AreEqual(
        1,
        fallback.CallCount);


    ProviderHealthRecord healthAfterFirstRun =
        _runtime.GetProviderHealth(
            primary);


    Assert.AreEqual(
        ProviderHealthState.CoolingDown,
        healthAfterFirstRun.State);


    Assert.AreEqual(
        2,
        healthAfterFirstRun.ConsecutiveFailures);


    Assert.IsFalse(
        _runtime.IsProviderAvailable(
            primary));


    int primaryCallsAfterFirstRun =
        primary.CallCount;


    //--------------------------------------------------
    // Second run:
    // primary must be skipped;
    // fallback succeeds again.
    //--------------------------------------------------

    Task<string> second =
        _runtime.ChatAsync(
            "second");


    yield return AwaitTask(
        second);


    Assert.AreEqual(
        "Fallback response",
        second.GetAwaiter()
            .GetResult());


    Assert.AreEqual(
        primaryCallsAfterFirstRun,
        primary.CallCount,
        "The primary provider should not be called while cooling down.");


    Assert.AreEqual(
        2,
        fallback.CallCount,
        "The fallback should handle both runs.");
}



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

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


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



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

        private sealed class FakeProvider :
            IProvider
        {
            private readonly int _failCount;


            public int CallCount
            {
                get;
                private set;
            }


            public string Name
            {
                get;
            }



            public FakeProvider(
                string name,
                int failCount)
            {
                Name =
                    name;


                _failCount =
                    failCount;
            }



            public Task<ChatCompletion> ChatAsync(
                IReadOnlyList<AIRuntime.Runtime.Models.ConversationMessage> messages,
                IReadOnlyList<ProviderTool> tools,
                CancellationToken cancellationToken)
            {
                CallCount++;


                if(CallCount <= _failCount)
                {
                    throw new TimeoutException(
                        "Simulated provider timeout.");
                }


                return Task.FromResult(
                    new ChatCompletion
                    {
                        Content =
                            Name +
                            " response"
                    });
            }
        }
    }
}
