using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;
using StellaOps.Auth.Client;
using StellaOps.Scanner.Core.Utility;

namespace StellaOps.Scanner.Core.Security;

public sealed class AuthorityTokenSource : IAuthorityTokenSource
{
    private readonly IStellaOpsTokenClient tokenClient;
    private readonly TimeProvider timeProvider;
    private readonly TimeSpan refreshSkew;
    private readonly ILogger<AuthorityTokenSource>? logger;
    private readonly ConcurrentDictionary<string, CacheEntry> cache = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, SemaphoreSlim> locks = new(StringComparer.Ordinal);

    public AuthorityTokenSource(
        IStellaOpsTokenClient tokenClient,
        TimeSpan? refreshSkew = null,
        TimeProvider? timeProvider = null,
        ILogger<AuthorityTokenSource>? logger = null)
    {
        this.tokenClient = tokenClient ?? throw new ArgumentNullException(nameof(tokenClient));
        this.timeProvider = timeProvider ?? TimeProvider.System;
        this.logger = logger;
        this.refreshSkew = refreshSkew is { } value && value > TimeSpan.Zero ? value : TimeSpan.FromSeconds(30);
    }

    public async ValueTask<ScannerOperationalToken> GetAsync(string audience, IEnumerable<string> scopes, CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(audience);

        var normalizedAudience = NormalizeAudience(audience);
        var normalizedScopes = NormalizeScopes(scopes, normalizedAudience);
        var cacheKey = BuildCacheKey(normalizedAudience, normalizedScopes);

        if (cache.TryGetValue(cacheKey, out var cached) && !cached.Token.IsExpired(timeProvider, refreshSkew))
        {
            return cached.Token;
        }

        var mutex = locks.GetOrAdd(cacheKey, static _ => new SemaphoreSlim(1, 1));
        await mutex.WaitAsync(cancellationToken).ConfigureAwait(false);

        try
        {
            if (cache.TryGetValue(cacheKey, out cached) && !cached.Token.IsExpired(timeProvider, refreshSkew))
            {
                return cached.Token;
            }

            var scopeString = string.Join(' ', normalizedScopes);
            var tokenResult = await tokenClient.RequestClientCredentialsTokenAsync(scopeString, cancellationToken).ConfigureAwait(false);

            var token = ScannerOperationalToken.FromResult(
                tokenResult.AccessToken,
                tokenResult.TokenType,
                tokenResult.ExpiresAtUtc,
                tokenResult.Scopes);

            cache[cacheKey] = new CacheEntry(token);
            logger?.LogDebug(
                "Issued new scanner OpTok for audience {Audience} with scopes {Scopes}; expires at {ExpiresAt}.",
                normalizedAudience,
                scopeString,
                token.ExpiresAt);

            return token;
        }
        finally
        {
            mutex.Release();
        }
    }

    public ValueTask InvalidateAsync(string audience, IEnumerable<string> scopes, CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(audience);

        var normalizedAudience = NormalizeAudience(audience);
        var normalizedScopes = NormalizeScopes(scopes, normalizedAudience);
        var cacheKey = BuildCacheKey(normalizedAudience, normalizedScopes);

        cache.TryRemove(cacheKey, out _);
        if (locks.TryRemove(cacheKey, out var mutex))
        {
            mutex.Dispose();
        }

        logger?.LogDebug("Invalidated cached OpTok for {Audience} ({CacheKey}).", normalizedAudience, cacheKey);
        return ValueTask.CompletedTask;
    }

    private static string NormalizeAudience(string audience)
        => audience.Trim().ToLowerInvariant();

    private static IReadOnlyList<string> NormalizeScopes(IEnumerable<string> scopes, string audience)
    {
        var set = new SortedSet<string>(StringComparer.Ordinal)
        {
            $"aud:{audience}"
        };

        if (scopes is not null)
        {
            foreach (var scope in scopes)
            {
                if (string.IsNullOrWhiteSpace(scope))
                {
                    continue;
                }

                set.Add(scope.Trim());
            }
        }

        return set.ToArray();
    }

    private static string BuildCacheKey(string audience, IReadOnlyList<string> scopes)
        => ScannerIdentifiers.CreateDeterministicHash(audience, string.Join(' ', scopes));

    private readonly record struct CacheEntry(ScannerOperationalToken Token);
}
