using System.Collections.Immutable;
using System.Formats.Asn1;
using System.Security.Cryptography;
using System.Security.Cryptography.Pkcs;
using System.Security.Cryptography.X509Certificates;
using System.Text;

namespace StellaOps.AirGap.Bundle.Services;

public interface ITsaChainBundler
{
    Task<TsaChainBundleResult> BundleAsync(
        ReadOnlyMemory<byte> timeStampToken,
        string outputPath,
        string? filePrefix = null,
        CancellationToken ct = default);
}

public sealed class TsaChainBundler : ITsaChainBundler
{
    public async Task<TsaChainBundleResult> BundleAsync(
        ReadOnlyMemory<byte> timeStampToken,
        string outputPath,
        string? filePrefix = null,
        CancellationToken ct = default)
    {
        if (timeStampToken.IsEmpty)
        {
            throw new ArgumentException("RFC3161 timestamp token is required.", nameof(timeStampToken));
        }

        var signedCms = new SignedCms();
        signedCms.Decode(timeStampToken.ToArray());

        if (signedCms.SignerInfos.Count == 0)
        {
            throw new InvalidOperationException("RFC3161 timestamp token has no signer.");
        }

        var signerCert = signedCms.SignerInfos[0].Certificate;
        if (signerCert is null)
        {
            throw new InvalidOperationException("RFC3161 timestamp token has no signer certificate.");
        }

        var certificates = new List<X509Certificate2>(signedCms.Certificates.Cast<X509Certificate2>());
        if (!certificates.Any(c => string.Equals(c.Thumbprint, signerCert.Thumbprint, StringComparison.OrdinalIgnoreCase)))
        {
            certificates.Add(signerCert);
        }

        var chain = BuildChain(signerCert, certificates);
        if (chain.Count == 0)
        {
            throw new InvalidOperationException("RFC3161 timestamp token contains no usable certificate chain.");
        }

        filePrefix ??= ComputePrefix(timeStampToken.Span);

        var entries = new List<TsaChainEntry>(chain.Count);
        for (var i = 0; i < chain.Count; i++)
        {
            var cert = chain[i];
            var certHash = ComputeShortHash(cert.RawData);
            var fileName = $"{filePrefix}-{i:D2}-{certHash}.pem";
            var relativePath = $"tsa/chain/{fileName}";
            var targetPath = PathValidation.SafeCombine(outputPath, relativePath);

            Directory.CreateDirectory(Path.GetDirectoryName(targetPath) ?? outputPath);
            await File.WriteAllTextAsync(targetPath, EncodePem(cert.RawData), Encoding.ASCII, ct)
                .ConfigureAwait(false);

            var info = new FileInfo(targetPath);
            entries.Add(new TsaChainEntry(cert, relativePath, info.Length));
        }

        return new TsaChainBundleResult(
            entries.Select(e => e.RelativePath).ToImmutableArray(),
            entries.Select(e => e.Certificate).ToImmutableArray(),
            entries.Sum(e => e.SizeBytes));
    }

    private static List<X509Certificate2> BuildChain(
        X509Certificate2 leaf,
        IReadOnlyList<X509Certificate2> pool)
    {
        var byThumbprint = new Dictionary<string, X509Certificate2>(StringComparer.OrdinalIgnoreCase);
        foreach (var cert in pool)
        {
            if (!string.IsNullOrWhiteSpace(cert.Thumbprint) && !byThumbprint.ContainsKey(cert.Thumbprint))
            {
                byThumbprint[cert.Thumbprint] = cert;
            }
        }

        var chain = new List<X509Certificate2>();
        var visited = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
        var current = leaf;

        while (current is not null && !string.IsNullOrWhiteSpace(current.Thumbprint))
        {
            if (!visited.Add(current.Thumbprint))
            {
                break;
            }

            chain.Add(current);

            if (IsSelfSigned(current))
            {
                break;
            }

            var issuer = FindIssuer(current, byThumbprint.Values);
            if (issuer is null)
            {
                break;
            }

            current = issuer;
        }

        return chain;
    }

    private static X509Certificate2? FindIssuer(
        X509Certificate2 certificate,
        IEnumerable<X509Certificate2> candidates)
    {
        var issuerName = certificate.Issuer;
        var issuerCandidates = candidates
            .Where(c => string.Equals(c.Subject, issuerName, StringComparison.OrdinalIgnoreCase))
            .OrderBy(c => c.Thumbprint, StringComparer.OrdinalIgnoreCase)
            .ToList();

        if (issuerCandidates.Count == 0)
        {
            return null;
        }

        if (issuerCandidates.Count == 1)
        {
            return issuerCandidates[0];
        }

        var authorityKeyId = TryGetAuthorityKeyIdentifier(certificate);
        if (authorityKeyId is null)
        {
            return issuerCandidates[0];
        }

        foreach (var candidate in issuerCandidates)
        {
            var subjectKeyId = TryGetSubjectKeyIdentifier(candidate);
            if (subjectKeyId is not null && subjectKeyId.SequenceEqual(authorityKeyId))
            {
                return candidate;
            }
        }

        return issuerCandidates[0];
    }

    private static bool IsSelfSigned(X509Certificate2 certificate)
    {
        if (!string.Equals(certificate.Subject, certificate.Issuer, StringComparison.OrdinalIgnoreCase))
        {
            return false;
        }

        var authorityKeyId = TryGetAuthorityKeyIdentifier(certificate);
        var subjectKeyId = TryGetSubjectKeyIdentifier(certificate);

        if (authorityKeyId is null || subjectKeyId is null)
        {
            return true;
        }

        return authorityKeyId.SequenceEqual(subjectKeyId);
    }

    private static byte[]? TryGetSubjectKeyIdentifier(X509Certificate2 certificate)
    {
        var ext = certificate.Extensions.Cast<X509Extension>()
            .FirstOrDefault(e => e.Oid?.Value == "2.5.29.14");
        if (ext is null)
        {
            return null;
        }

        try
        {
            var ski = new X509SubjectKeyIdentifierExtension(ext, ext.Critical);
            var keyId = ski.SubjectKeyIdentifier;
            if (string.IsNullOrWhiteSpace(keyId))
            {
                return null;
            }

            return Convert.FromHexString(keyId);
        }
        catch
        {
            return null;
        }
    }

    private static byte[]? TryGetAuthorityKeyIdentifier(X509Certificate2 certificate)
    {
        var ext = certificate.Extensions.Cast<X509Extension>()
            .FirstOrDefault(e => e.Oid?.Value == "2.5.29.35");
        if (ext is null)
        {
            return null;
        }

        try
        {
            var reader = new AsnReader(ext.RawData, AsnEncodingRules.DER);
            var akiBytes = reader.ReadOctetString();
            var akiReader = new AsnReader(akiBytes, AsnEncodingRules.DER);
            var sequence = akiReader.ReadSequence();

            while (sequence.HasData)
            {
                var tag = sequence.PeekTag();
                if (tag.TagClass == TagClass.ContextSpecific && tag.TagValue == 0)
                {
                    return sequence.ReadOctetString(new Asn1Tag(TagClass.ContextSpecific, 0));
                }

                sequence.ReadEncodedValue();
            }
        }
        catch
        {
            return null;
        }

        return null;
    }

    private static string ComputePrefix(ReadOnlySpan<byte> tokenBytes)
    {
        var hash = SHA256.HashData(tokenBytes);
        return Convert.ToHexString(hash).ToLowerInvariant()[..12];
    }

    private static string ComputeShortHash(byte[] data)
    {
        var hash = SHA256.HashData(data);
        return Convert.ToHexString(hash).ToLowerInvariant()[..16];
    }

    private static string EncodePem(byte[] raw)
    {
        var base64 = Convert.ToBase64String(raw, Base64FormattingOptions.InsertLineBreaks);
        var builder = new StringBuilder();
        builder.Append("-----BEGIN CERTIFICATE-----\n");
        builder.Append(base64);
        builder.Append("\n-----END CERTIFICATE-----\n");
        return builder.ToString();
    }
}

public sealed record TsaChainBundleResult(
    ImmutableArray<string> ChainPaths,
    ImmutableArray<X509Certificate2> Certificates,
    long TotalSizeBytes);

internal sealed record TsaChainEntry(X509Certificate2 Certificate, string RelativePath, long SizeBytes);