using ZB.MOM.WW.Secrets.Abstractions;

namespace ZB.MOM.WW.Secrets.Rotation;

/// <summary>
/// Rotates the master key-encryption key (KEK) across every stored secret by re-wrapping each
/// row's data-encryption key (DEK) from the old KEK to the new one. Because the envelope model
/// wraps a per-secret DEK under the KEK, rotation never re-encrypts secret bodies — only the DEK
/// wrap changes — so plaintext is never exposed and no value history is touched.
/// </summary>
/// <remarks>
/// <para>
/// The pass is <b>idempotent and resumable</b>: a row already wrapped by the new KEK is skipped
/// (counted in <see cref="RewrapReport.AlreadyCurrent"/>), so a run interrupted part-way can simply
/// be re-run to finish. It is <b>fail-closed</b>: a row wrapped by neither the old nor the new KEK
/// aborts the pass (rather than silently leaving an un-migratable row that would become
/// undecryptable once the old KEK is retired); rows already re-wrapped in the aborted run stay
/// persisted, so a re-run resumes cleanly after the anomaly is investigated.
/// </para>
/// <para>
/// Rotation is a per-store operation and must be run once per independent store (once for a shared
/// SQL-Server store; once per node for a per-node SQLite store on a shared KEK). It should be run
/// with resolve traffic quiesced: the persistence path deliberately does not bump revision or
/// <c>updated_utc</c>, so a re-wrap is invisible to cluster last-writer-wins reconciliation.
/// </para>
/// </remarks>
public sealed class KekRotationService
{
    private readonly ISecretStore _store;
    private readonly ISecretCipher _cipher;

    /// <summary>Creates the rotation service over the encrypted store and envelope cipher.</summary>
    /// <param name="store">The store to enumerate and re-wrap in place.</param>
    /// <param name="cipher">The envelope cipher providing the DEK re-wrap primitive.</param>
    public KekRotationService(ISecretStore store, ISecretCipher cipher)
    {
        _store = store ?? throw new ArgumentNullException(nameof(store));
        _cipher = cipher ?? throw new ArgumentNullException(nameof(cipher));
    }

    /// <summary>
    /// Re-wraps every row from <paramref name="oldKek"/> onto <paramref name="newKek"/>.
    /// </summary>
    /// <param name="oldKek">The provider for the KEK rows are currently wrapped under.</param>
    /// <param name="newKek">The provider for the KEK rows should be re-wrapped onto.</param>
    /// <param name="ct">A token to cancel the pass (already-persisted re-wraps are retained).</param>
    /// <returns>A <see cref="RewrapReport"/> summarizing the pass (no secret material).</returns>
    /// <exception cref="ArgumentException">
    /// The old and new KEK identifiers are the same — there is nothing to rotate.
    /// </exception>
    /// <exception cref="SecretDecryptionException">
    /// A row is wrapped by a KEK that is neither the old nor the new one (the pass aborts), or a
    /// row's DEK fails to unwrap under the supplied old KEK.
    /// </exception>
    public async Task<RewrapReport> RewrapAllAsync(
        IMasterKeyProvider oldKek, IMasterKeyProvider newKek, CancellationToken ct)
    {
        ArgumentNullException.ThrowIfNull(oldKek);
        ArgumentNullException.ThrowIfNull(newKek);

        if (string.Equals(oldKek.KekId, newKek.KekId, StringComparison.Ordinal))
        {
            throw new ArgumentException(
                $"The old and new KEK identifiers are identical ('{oldKek.KekId}'); there is nothing to rotate.",
                nameof(newKek));
        }

        // Enumerate metadata (never ciphertext) for ALL rows including tombstones — a tombstoned row
        // still carries a KEK-wrapped DEK and would become un-migratable if left on a retired KEK.
        IReadOnlyList<SecretMetadata> all =
            await _store.ListAsync(includeDeleted: true, ct).ConfigureAwait(false);

        int rewrapped = 0;
        int alreadyCurrent = 0;

        foreach (SecretMetadata meta in all)
        {
            ct.ThrowIfCancellationRequested();

            // Fast-path skip on the metadata kek_id (no ciphertext fetch): a row already on the new
            // KEK is a completed migration — this is what makes the pass idempotent / resumable.
            if (string.Equals(meta.KekId, newKek.KekId, StringComparison.Ordinal))
            {
                alreadyCurrent++;
                continue;
            }

            switch (await RewrapOneAsync(meta.Name, oldKek, newKek, rewrapped, ct).ConfigureAwait(false))
            {
                case RowOutcome.Rewrapped:
                    rewrapped++;
                    break;
                case RowOutcome.AlreadyCurrent:
                    alreadyCurrent++;
                    break;
                case RowOutcome.Vanished:
                    break;
            }
        }

        return new RewrapReport { Total = all.Count, Rewrapped = rewrapped, AlreadyCurrent = alreadyCurrent };
    }

    private enum RowOutcome { Rewrapped, AlreadyCurrent, Vanished }

    // Re-wraps a single row with a bounded fetch→rewrap→compare-and-swap retry so a rare concurrent
    // write (the pass should run with writes quiesced) is retried rather than corrupting the row.
    private async Task<RowOutcome> RewrapOneAsync(
        SecretName name, IMasterKeyProvider oldKek, IMasterKeyProvider newKek, int rewrappedSoFar, CancellationToken ct)
    {
        const int maxAttempts = 5;
        for (int attempt = 1; ; attempt++)
        {
            ct.ThrowIfCancellationRequested();

            StoredSecret? row = await _store.GetAsync(name, ct).ConfigureAwait(false);
            if (row is null)
            {
                // Vanished between the list and the fetch (a concurrent hard-absence).
                return RowOutcome.Vanished;
            }

            // Re-evaluate against the FRESH row: a concurrent write may have already moved it.
            if (string.Equals(row.KekId, newKek.KekId, StringComparison.Ordinal))
            {
                return RowOutcome.AlreadyCurrent;
            }

            if (!string.Equals(row.KekId, oldKek.KekId, StringComparison.Ordinal))
            {
                throw new SecretDecryptionException(
                    $"Secret '{name.Value}' is wrapped by KEK '{row.KekId}', which is neither the old " +
                    $"('{oldKek.KekId}') nor the new ('{newKek.KekId}') KEK. Rotation aborted after re-wrapping " +
                    $"{rewrappedSoFar} row(s); investigate the anomalous row and re-run to resume.");
            }

            // Rewrap re-verifies row.KekId == oldKek.KekId and fails closed on a wrong old key.
            StoredSecret rewrappedRow = _cipher.Rewrap(row, oldKek, newKek);

            // Compare-and-swap on the wrap we just unwrapped: a concurrent set/rotate changed it →
            // 0 rows → retry (fetch the newer row and re-evaluate).
            if (await _store.ApplyRewrapAsync(rewrappedRow, row.WrappedDek, ct).ConfigureAwait(false))
            {
                return RowOutcome.Rewrapped;
            }

            if (attempt >= maxAttempts)
            {
                throw new SecretDecryptionException(
                    $"Secret '{name.Value}' is being modified concurrently; rotation could not converge after " +
                    $"{maxAttempts} attempts. Re-run with secret writes quiesced.");
            }
        }
    }
}
