using System;
using System.Collections.Generic;
using System.Diagnostics;

using AIRuntime.Runtime.Providers.Fallback;


namespace AIRuntime.Runtime.Providers.Health
{
    /// <summary>
    /// Tracks provider success, failure, latency and cooldown state.
    /// </summary>
    public sealed class ProviderHealthTracker
    {
        private readonly Dictionary<IProvider, ProviderHealthRecord> _records =
            new Dictionary<IProvider, ProviderHealthRecord>();


        private readonly ProviderCooldownPolicy _cooldownPolicy;


        public ProviderHealthTracker(
            ProviderCooldownPolicy cooldownPolicy)
        {
            _cooldownPolicy =
                cooldownPolicy ??
                throw new ArgumentNullException(
                    nameof(cooldownPolicy));
        }


        public ProviderHealthRecord GetHealth(
            IProvider provider)
        {
            if (provider == null)
            {
                throw new ArgumentNullException(
                    nameof(provider));
            }


            ProviderHealthRecord record =
                GetOrCreateRecord(
                    provider);


            RefreshState(
                record,
                DateTimeOffset.UtcNow);


            return record.CreateSnapshot();
        }


        public IReadOnlyList<ProviderHealthRecord> GetAllHealth()
        {
            DateTimeOffset now =
                DateTimeOffset.UtcNow;


            List<ProviderHealthRecord> results =
                new List<ProviderHealthRecord>(
                    _records.Count);


            foreach (
                ProviderHealthRecord record
                in _records.Values)
            {
                RefreshState(
                    record,
                    now);


                results.Add(
                    record.CreateSnapshot());
            }


            return results;
        }


        public bool IsAvailable(
            IProvider provider)
        {
            ProviderHealthRecord record =
                GetOrCreateRecord(
                    provider);


            DateTimeOffset now =
                DateTimeOffset.UtcNow;


            RefreshState(
                record,
                now);


            return !record.IsCoolingDown(
                now);
        }


        public void RecordAttempt(
            IProvider provider)
        {
            ProviderHealthRecord record =
                GetOrCreateRecord(
                    provider);


            record.TotalAttempts++;


            record.LastAttemptAt =
                DateTimeOffset.UtcNow;
        }


        public void RecordSuccess(
            IProvider provider,
            TimeSpan latency)
        {
            ProviderHealthRecord record =
                GetOrCreateRecord(
                    provider);


            record.TotalSuccesses++;


            record.ConsecutiveFailures =
                0;


            record.LastSuccessAt =
                DateTimeOffset.UtcNow;


            record.LastLatency =
                latency;


            record.CooldownUntil =
                null;


            record.State =
                ProviderHealthState.Healthy;
        }


        public void RecordFailure(
            IProvider provider,
            ProviderFailure failure,
            TimeSpan latency)
        {
            if (failure == null)
            {
                throw new ArgumentNullException(
                    nameof(failure));
            }


            ProviderHealthRecord record =
                GetOrCreateRecord(
                    provider);


            record.TotalFailures++;


            record.LastFailureAt =
                DateTimeOffset.UtcNow;


            record.LastFailureKind =
                failure.Kind;


            record.LastFailureMessage =
                failure.Message;


            record.LastLatency =
                latency;


            if (
                !_cooldownPolicy.CountsTowardsCooldown(
                    failure.Kind))
            {
                record.State =
                    ProviderHealthState.Degraded;

                return;
            }


            record.ConsecutiveFailures++;


            TimeSpan cooldown =
                _cooldownPolicy.GetCooldown(
                    record.ConsecutiveFailures,
                    failure.Kind);


            if (cooldown > TimeSpan.Zero)
            {
                record.CooldownUntil =
                    DateTimeOffset.UtcNow +
                    cooldown;


                record.State =
                    ProviderHealthState.CoolingDown;
            }
            else
            {
                record.State =
                    ProviderHealthState.Degraded;
            }
        }


        public void Reset(
            IProvider provider)
        {
            if (provider == null)
                return;


            _records.Remove(
                provider);
        }


        public void ResetAll()
        {
            _records.Clear();
        }


        private ProviderHealthRecord GetOrCreateRecord(
            IProvider provider)
        {
            if (provider == null)
            {
                throw new ArgumentNullException(
                    nameof(provider));
            }


            if (
                _records.TryGetValue(
                    provider,
                    out ProviderHealthRecord record))
            {
                return record;
            }


            record =
                new ProviderHealthRecord(
                    provider);


            _records.Add(
                provider,
                record);


            return record;
        }


        private static void RefreshState(
            ProviderHealthRecord record,
            DateTimeOffset now)
        {
            if (!record.CooldownUntil.HasValue)
                return;


            if (record.CooldownUntil.Value > now)
            {
                record.State =
                    ProviderHealthState.CoolingDown;

                return;
            }


            record.CooldownUntil =
                null;


            record.State =
                record.ConsecutiveFailures > 0
                    ? ProviderHealthState.Degraded
                    : ProviderHealthState.Healthy;
        }
    }
}