using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Globalization;
using System.Linq;
using System.Net;
using System.Net.Http;
using System.Net.Http.Headers;
using System.Text;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;
using StellaOps.Excititor.Connectors.OCI.OpenVEX.Attest.Authentication;
using StellaOps.Excititor.Connectors.OCI.OpenVEX.Attest.Configuration;
using StellaOps.Excititor.Connectors.OCI.OpenVEX.Attest.Discovery;

namespace StellaOps.Excititor.Connectors.OCI.OpenVEX.Attest.Fetch;

internal sealed class OciRegistryClient
{
    private const string ManifestMediaType = "application/vnd.oci.image.manifest.v1+json";
    private const string ReferrersArtifactType = "application/vnd.dsse.envelope.v1+json";
    private const string DsseMediaType = "application/vnd.dsse.envelope.v1+json";
    private const string OpenVexMediaType = "application/vnd.cncf.openvex.v1+json";

    private static readonly JsonSerializerOptions SerializerOptions = new(JsonSerializerDefaults.Web);

    private readonly IHttpClientFactory _httpClientFactory;
    private readonly ILogger _logger;
    private readonly OciRegistryAuthorization _authorization;
    private readonly OciOpenVexAttestationConnectorOptions _options;

    public OciRegistryClient(
        IHttpClientFactory httpClientFactory,
        ILogger logger,
        OciRegistryAuthorization authorization,
        OciOpenVexAttestationConnectorOptions options)
    {
        _httpClientFactory = httpClientFactory ?? throw new ArgumentNullException(nameof(httpClientFactory));
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
        _authorization = authorization ?? throw new ArgumentNullException(nameof(authorization));
        _options = options ?? throw new ArgumentNullException(nameof(options));
    }

    public async Task<string?> ResolveDigestAsync(OciImageReference image, CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(image);

        if (image.HasDigest)
        {
            return image.Digest;
        }

        var requestUri = BuildRegistryUri(image, $"manifests/{EscapeReference(image.Tag ?? "latest")}");

        async Task<HttpRequestMessage> RequestFactory()
        {
            var request = new HttpRequestMessage(HttpMethod.Head, requestUri);
            request.Headers.Accept.ParseAdd(ManifestMediaType);
            ApplyAuthentication(request);
            return await Task.FromResult(request).ConfigureAwait(false);
        }

        using var response = await SendAsync(RequestFactory, cancellationToken).ConfigureAwait(false);

        if (!response.IsSuccessStatusCode)
        {
            if (response.StatusCode == HttpStatusCode.NotFound)
            {
                _logger.LogWarning("Failed to resolve digest for {Reference}; registry returned 404.", image.Canonical);
                return null;
            }

            response.EnsureSuccessStatusCode();
        }

        if (response.Headers.TryGetValues("Docker-Content-Digest", out var values))
        {
            var digest = values.FirstOrDefault();
            if (!string.IsNullOrWhiteSpace(digest))
            {
                return digest;
            }
        }

        // Manifest may have been returned without digest header; fall back to GET.
        async Task<HttpRequestMessage> ManifestRequestFactory()
        {
            var request = new HttpRequestMessage(HttpMethod.Get, requestUri);
            request.Headers.Accept.ParseAdd(ManifestMediaType);
            ApplyAuthentication(request);
            return await Task.FromResult(request).ConfigureAwait(false);
        }

        using var manifestResponse = await SendAsync(ManifestRequestFactory, cancellationToken).ConfigureAwait(false);
        manifestResponse.EnsureSuccessStatusCode();

        if (manifestResponse.Headers.TryGetValues("Docker-Content-Digest", out var manifestValues))
        {
            return manifestValues.FirstOrDefault();
        }

        _logger.LogWarning("Registry {Registry} did not provide Docker-Content-Digest header for {Reference}.", image.Registry, image.Canonical);
        return null;
    }

    public async Task<IReadOnlyList<OciArtifactDescriptor>> ListReferrersAsync(
        OciImageReference image,
        string subjectDigest,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(image);
        ArgumentNullException.ThrowIfNull(subjectDigest);

        var query = $"artifactType={Uri.EscapeDataString(ReferrersArtifactType)}";
        var requestUri = BuildRegistryUri(image, $"referrers/{subjectDigest}", query);

        async Task<HttpRequestMessage> RequestFactory()
        {
            var request = new HttpRequestMessage(HttpMethod.Get, requestUri);
            ApplyAuthentication(request);
            request.Headers.Accept.ParseAdd("application/json");
            return await Task.FromResult(request).ConfigureAwait(false);
        }

        using var response = await SendAsync(RequestFactory, cancellationToken).ConfigureAwait(false);
        if (!response.IsSuccessStatusCode)
        {
            if (response.StatusCode == HttpStatusCode.NotFound)
            {
                _logger.LogDebug("Registry returned 404 for referrers on {Subject}.", subjectDigest);
                return Array.Empty<OciArtifactDescriptor>();
            }

            response.EnsureSuccessStatusCode();
        }

        await using var stream = await response.Content.ReadAsStreamAsync(cancellationToken).ConfigureAwait(false);
        var index = await JsonSerializer.DeserializeAsync<OciReferrerIndex>(stream, SerializerOptions, cancellationToken).ConfigureAwait(false);
        return index?.Referrers ?? Array.Empty<OciArtifactDescriptor>();
    }

    public async Task<OciAttestationDocument?> DownloadAttestationAsync(
        OciImageReference image,
        OciArtifactDescriptor descriptor,
        string subjectDigest,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(image);
        ArgumentNullException.ThrowIfNull(descriptor);

        if (!IsSupportedDescriptor(descriptor))
        {
            return null;
        }

        var requestUri = BuildRegistryUri(image, $"blobs/{descriptor.Digest}");

        async Task<HttpRequestMessage> RequestFactory()
        {
            var request = new HttpRequestMessage(HttpMethod.Get, requestUri);
            ApplyAuthentication(request);
            request.Headers.Accept.ParseAdd(descriptor.MediaType ?? "application/octet-stream");
            return await Task.FromResult(request).ConfigureAwait(false);
        }

        using var response = await SendAsync(RequestFactory, cancellationToken).ConfigureAwait(false);
        if (!response.IsSuccessStatusCode)
        {
            if (response.StatusCode == HttpStatusCode.NotFound)
            {
                _logger.LogWarning("Registry returned 404 while downloading attestation {Digest} for {Subject}.", descriptor.Digest, subjectDigest);
                return null;
            }

            response.EnsureSuccessStatusCode();
        }

        var buffer = await response.Content.ReadAsByteArrayAsync(cancellationToken).ConfigureAwait(false);
        var metadata = BuildMetadata(image, descriptor, "registry", requestUri.ToString(), subjectDigest);
        return new OciAttestationDocument(
            requestUri,
            buffer,
            metadata,
            subjectDigest,
            descriptor.Digest,
            descriptor.ArtifactType,
            "registry");
    }

    private static bool IsSupportedDescriptor(OciArtifactDescriptor descriptor)
    {
        if (descriptor is null)
        {
            return false;
        }

        if (!string.IsNullOrWhiteSpace(descriptor.ArtifactType) &&
            descriptor.ArtifactType.Equals(OpenVexMediaType, StringComparison.OrdinalIgnoreCase))
        {
            return true;
        }

        if (!string.IsNullOrWhiteSpace(descriptor.MediaType) &&
            (descriptor.MediaType.Equals(DsseMediaType, StringComparison.OrdinalIgnoreCase) ||
             descriptor.MediaType.Equals(OpenVexMediaType, StringComparison.OrdinalIgnoreCase)))
        {
            return true;
        }

        return false;
    }

    private async Task<HttpResponseMessage> SendAsync(
        Func<Task<HttpRequestMessage>> requestFactory,
        CancellationToken cancellationToken)
    {
        const int maxAttempts = 3;
        TimeSpan delay = TimeSpan.FromSeconds(1);
        Exception? lastError = null;

        for (var attempt = 1; attempt <= maxAttempts; attempt++)
        {
            cancellationToken.ThrowIfCancellationRequested();

            using var request = await requestFactory().ConfigureAwait(false);
            var client = _httpClientFactory.CreateClient(OciOpenVexAttestationConnectorOptions.HttpClientName);

            try
            {
                var response = await client.SendAsync(request, cancellationToken).ConfigureAwait(false);

                if (response.IsSuccessStatusCode)
                {
                    return response;
                }

                if (response.StatusCode == HttpStatusCode.Unauthorized)
                {
                    if (_authorization.Mode == OciRegistryAuthMode.Anonymous && !_authorization.AllowAnonymousFallback)
                    {
                        var message = $"Registry request to {request.RequestUri} was unauthorized and anonymous fallback is disabled.";
                        response.Dispose();
                        throw new HttpRequestException(message);
                    }

                    lastError = new HttpRequestException($"Registry returned 401 Unauthorized for {request.RequestUri}.");
                }
                else if ((int)response.StatusCode >= 500 || response.StatusCode == (HttpStatusCode)429)
                {
                    lastError = new HttpRequestException($"Registry returned status {(int)response.StatusCode} ({response.ReasonPhrase}) for {request.RequestUri}.");
                }
                else
                {
                    response.EnsureSuccessStatusCode();
                }

                response.Dispose();
            }
            catch (Exception ex) when (ex is HttpRequestException or TaskCanceledException)
            {
                lastError = ex;
            }

            if (attempt < maxAttempts)
            {
                await Task.Delay(delay, cancellationToken).ConfigureAwait(false);
                delay = TimeSpan.FromSeconds(Math.Min(delay.TotalSeconds * 2, 10));
            }
        }

        throw new HttpRequestException("Failed to execute OCI registry request after multiple attempts.", lastError);
    }

    private void ApplyAuthentication(HttpRequestMessage request)
    {
        switch (_authorization.Mode)
        {
            case OciRegistryAuthMode.Basic when
                !string.IsNullOrEmpty(_authorization.Username) &&
                !string.IsNullOrEmpty(_authorization.Password):
                var credentials = Convert.ToBase64String(Encoding.UTF8.GetBytes($"{_authorization.Username}:{_authorization.Password}"));
                request.Headers.Authorization = new AuthenticationHeaderValue("Basic", credentials);
                break;
            case OciRegistryAuthMode.IdentityToken when !string.IsNullOrWhiteSpace(_authorization.IdentityToken):
                request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _authorization.IdentityToken);
                break;
            case OciRegistryAuthMode.RefreshToken when !string.IsNullOrWhiteSpace(_authorization.RefreshToken):
                request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _authorization.RefreshToken);
                break;
            default:
                if (_authorization.Mode != OciRegistryAuthMode.Anonymous && !_authorization.AllowAnonymousFallback)
                {
                    _logger.LogDebug("No authentication header applied for request to {Uri} (mode {Mode}).", request.RequestUri, _authorization.Mode);
                }
                break;
        }
    }

    private Uri BuildRegistryUri(OciImageReference image, string relativePath, string? query = null)
    {
        var scheme = image.Scheme;
        if (!string.Equals(scheme, "https", StringComparison.OrdinalIgnoreCase) && !_options.AllowHttpRegistries)
        {
            throw new InvalidOperationException($"HTTP access to registry '{image.Registry}' is disabled. Set AllowHttpRegistries to true to enable.");
        }

        var builder = new UriBuilder($"{scheme}://{image.Registry}")
        {
            Path = $"v2/{BuildRepositoryPath(image.Repository)}/{relativePath}"
        };

        if (!string.IsNullOrWhiteSpace(query))
        {
            builder.Query = query;
        }

        return builder.Uri;
    }

    private static string BuildRepositoryPath(string repository)
    {
        var segments = repository.Split('/', StringSplitOptions.RemoveEmptyEntries);
        return string.Join('/', segments.Select(Uri.EscapeDataString));
    }

    private static string EscapeReference(string reference)
    {
        return Uri.EscapeDataString(reference);
    }

    private static ImmutableDictionary<string, string> BuildMetadata(
        OciImageReference image,
        OciArtifactDescriptor descriptor,
        string sourceKind,
        string sourcePath,
        string subjectDigest)
    {
        var builder = ImmutableDictionary.CreateBuilder<string, string>(StringComparer.Ordinal);
        builder["oci.image.registry"] = image.Registry;
        builder["oci.image.repository"] = image.Repository;
        builder["oci.image.reference"] = image.Canonical;
        builder["oci.image.subjectDigest"] = subjectDigest;
        builder["oci.attestation.sourceKind"] = sourceKind;
        builder["oci.attestation.source"] = sourcePath;
        builder["oci.attestation.artifactDigest"] = descriptor.Digest;
        builder["oci.attestation.mediaType"] = descriptor.MediaType ?? string.Empty;
        builder["oci.attestation.artifactType"] = descriptor.ArtifactType ?? string.Empty;
        builder["oci.attestation.size"] = descriptor.Size.ToString(CultureInfo.InvariantCulture);

        if (descriptor.Annotations is not null)
        {
            foreach (var annotation in descriptor.Annotations)
            {
                builder[$"oci.attestation.annotations.{annotation.Key}"] = annotation.Value;
            }
        }

        return builder.ToImmutable();
    }
}
