using System.Buffers;
using System.Security.Cryptography;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using StellaOps.Scanner.Storage.Catalog;
using StellaOps.Scanner.Storage.ObjectStore;
using StellaOps.Scanner.Storage.Repositories;

namespace StellaOps.Scanner.Storage.Services;

public sealed class ArtifactStorageService
{
    private readonly ArtifactRepository _artifactRepository;
    private readonly LifecycleRuleRepository _lifecycleRuleRepository;
    private readonly IArtifactObjectStore _objectStore;
    private readonly ScannerStorageOptions _options;
    private readonly ILogger<ArtifactStorageService> _logger;

    public ArtifactStorageService(
        ArtifactRepository artifactRepository,
        LifecycleRuleRepository lifecycleRuleRepository,
        IArtifactObjectStore objectStore,
        IOptions<ScannerStorageOptions> options,
        ILogger<ArtifactStorageService> logger)
    {
        _artifactRepository = artifactRepository ?? throw new ArgumentNullException(nameof(artifactRepository));
        _lifecycleRuleRepository = lifecycleRuleRepository ?? throw new ArgumentNullException(nameof(lifecycleRuleRepository));
        _objectStore = objectStore ?? throw new ArgumentNullException(nameof(objectStore));
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
        _options = (options ?? throw new ArgumentNullException(nameof(options))).Value;
    }

    public async Task<ArtifactDocument> StoreArtifactAsync(
        ArtifactDocumentType type,
        ArtifactDocumentFormat format,
        string mediaType,
        Stream content,
        bool immutable,
        string ttlClass,
        DateTime? expiresAtUtc,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(content);
        ArgumentException.ThrowIfNullOrWhiteSpace(mediaType);

        var (buffer, size, digestHex) = await BufferAndHashAsync(content, cancellationToken).ConfigureAwait(false);
        try
        {
            var normalizedDigest = $"sha256:{digestHex}";
            var artifactId = CatalogIdFactory.CreateArtifactId(type, normalizedDigest);
            var key = BuildObjectKey(type, format, normalizedDigest);
            var descriptor = new ArtifactObjectDescriptor(
                _options.ObjectStore.BucketName,
                key,
                immutable,
                _options.ObjectStore.ComplianceRetention);

            buffer.Position = 0;
            await _objectStore.PutAsync(descriptor, buffer, cancellationToken).ConfigureAwait(false);

            if (_options.DualWrite.Enabled)
            {
                buffer.Position = 0;
                var mirrorDescriptor = descriptor with { Bucket = _options.DualWrite.MirrorBucket! };
                await _objectStore.PutAsync(mirrorDescriptor, buffer, cancellationToken).ConfigureAwait(false);
            }

            var document = new ArtifactDocument
            {
                Id = artifactId,
                Type = type,
                Format = format,
                MediaType = mediaType,
                BytesSha256 = normalizedDigest,
                SizeBytes = size,
                Immutable = immutable,
                RefCount = 1,
                CreatedAtUtc = DateTime.UtcNow,
                UpdatedAtUtc = DateTime.UtcNow,
                TtlClass = ttlClass,
            };

            await _artifactRepository.UpsertAsync(document, cancellationToken).ConfigureAwait(false);

            if (expiresAtUtc.HasValue)
            {
                var lifecycle = new LifecycleRuleDocument
                {
                    Id = CatalogIdFactory.CreateLifecycleRuleId(document.Id, ttlClass),
                    ArtifactId = document.Id,
                    Class = ttlClass,
                    ExpiresAtUtc = expiresAtUtc,
                    CreatedAtUtc = DateTime.UtcNow,
                };

                await _lifecycleRuleRepository.UpsertAsync(lifecycle, cancellationToken).ConfigureAwait(false);
            }

            _logger.LogInformation("Stored scanner artifact {ArtifactId} ({SizeBytes} bytes, digest {Digest})", document.Id, size, normalizedDigest);
            return document;
        }
        finally
        {
            await buffer.DisposeAsync().ConfigureAwait(false);
        }
    }

    private static async Task<(MemoryStream Buffer, long Size, string DigestHex)> BufferAndHashAsync(Stream content, CancellationToken cancellationToken)
    {
        var bufferStream = new MemoryStream();
        var hasher = IncrementalHash.CreateHash(HashAlgorithmName.SHA256);
        var rented = ArrayPool<byte>.Shared.Rent(81920);
        long total = 0;

        try
        {
            int read;
            while ((read = await content.ReadAsync(rented.AsMemory(0, rented.Length), cancellationToken).ConfigureAwait(false)) > 0)
            {
                hasher.AppendData(rented, 0, read);
                await bufferStream.WriteAsync(rented.AsMemory(0, read), cancellationToken).ConfigureAwait(false);
                total += read;
            }
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(rented);
        }

        bufferStream.Position = 0;
        var digest = hasher.GetCurrentHash();
        var digestHex = Convert.ToHexString(digest).ToLowerInvariant();
        return (bufferStream, total, digestHex);
    }

    private string BuildObjectKey(ArtifactDocumentType type, ArtifactDocumentFormat format, string digest)
    {
        var normalizedDigest = digest.Split(':', 2, StringSplitOptions.TrimEntries)[^1];
        var prefix = type switch
        {
            ArtifactDocumentType.LayerBom => ScannerStorageDefaults.ObjectPrefixes.Layers,
            ArtifactDocumentType.ImageBom => ScannerStorageDefaults.ObjectPrefixes.Images,
            ArtifactDocumentType.Diff => "diffs",
            ArtifactDocumentType.Index => ScannerStorageDefaults.ObjectPrefixes.Indexes,
            ArtifactDocumentType.Attestation => ScannerStorageDefaults.ObjectPrefixes.Attestations,
            _ => ScannerStorageDefaults.ObjectPrefixes.Images,
        };

        var extension = format switch
        {
            ArtifactDocumentFormat.CycloneDxJson => "sbom.cdx.json",
            ArtifactDocumentFormat.CycloneDxProtobuf => "sbom.cdx.pb",
            ArtifactDocumentFormat.SpdxJson => "sbom.spdx.json",
            ArtifactDocumentFormat.BomIndex => "bom-index.bin",
            ArtifactDocumentFormat.DsseJson => "artifact.dsse.json",
            _ => "artifact.bin",
        };

        var rootPrefix = _options.ObjectStore.RootPrefix;
        if (string.IsNullOrWhiteSpace(rootPrefix))
        {
            return $"{prefix}/{normalizedDigest}/{extension}";
        }

        return $"{TrimTrailingSlash(rootPrefix)}/{prefix}/{normalizedDigest}/{extension}";
    }

    private static string TrimTrailingSlash(string prefix)
    {
        if (string.IsNullOrWhiteSpace(prefix))
        {
            return string.Empty;
        }

        return prefix.TrimEnd('/');
    }
}
