using System.Diagnostics.CodeAnalysis;
using System.Security.Cryptography;
using DodoSSH.Client.Storage;
using DodoSSH.Crypto;

namespace DodoSSH.Client.Sync;

/// <summary>
/// The vault keys held for the duration of an unlocked session.
/// </summary>
/// <remarks>
/// <para>
/// One place that holds plaintext vault keys, so there is one place that clears them. Every store and
/// every cipher call borrows a key from here rather than keeping a copy, which is what makes
/// "the keys exist only while unlocked" a property of the code and not of everyone's discipline.
/// </para>
/// <para>
/// A grant that will not open is not an error: it means the vault has been rekeyed and this client's
/// grant has not been re-wrapped yet, or the grant was fabricated. Both leave the vault temporarily
/// unreadable and both are reported rather than thrown, so one bad grant does not take the other vaults
/// down with it.
/// </para>
/// </remarks>
public sealed class VaultKeyring : IDisposable
{
    private readonly Dictionary<Guid, byte[]> keys = [];
    private readonly Dictionary<Guid, uint> generations = [];
    private bool disposed;

    private VaultKeyring()
    {
    }

    /// <summary>Vaults whose grant could not be opened, and which are therefore unreadable.</summary>
    public IReadOnlyList<Guid> Unopened { get; private set; } = [];

    /// <summary>
    /// Opens every grant the bundle can.
    /// </summary>
    /// <param name="bundle">The unlocked identity keys.</param>
    /// <param name="vaults">The cached vault list, each with its wrapped key.</param>
    public static VaultKeyring Open(UserSecretBundle bundle, IReadOnlyList<StoredVault> vaults)
    {
        ArgumentNullException.ThrowIfNull(bundle);
        ArgumentNullException.ThrowIfNull(vaults);

        var keyring = new VaultKeyring();
        var unopened = new List<Guid>();

        try
        {
            foreach (var vault in vaults)
            {
                if (vault.WrappedVaultKey is null)
                {
                    // The server said so itself: a grant awaiting re-wrap after a rekey.
                    unopened.Add(vault.VaultId);
                    continue;
                }

                var key = VaultKeys.TryUnwrap(
                    bundle.EncryptionKey,
                    vault.WrappedVaultKey,
                    vault.VaultId,
                    vault.KeyGeneration);

                if (key is null)
                {
                    unopened.Add(vault.VaultId);
                    continue;
                }

                keyring.keys[vault.VaultId] = key;
                keyring.generations[vault.VaultId] = vault.KeyGeneration;
            }

            keyring.Unopened = unopened;
            return keyring;
        }
        catch
        {
            keyring.Dispose();
            throw;
        }
    }

    /// <summary>
    /// Takes a vault key this session has just generated.
    /// </summary>
    /// <param name="vaultId">The vault.</param>
    /// <param name="vaultKey">
    /// The plaintext key. <b>The keyring takes ownership</b> and zeroes it on disposal; the caller must
    /// not keep a reference or zero it itself.
    /// </param>
    /// <param name="keyGeneration">The generation this key is for.</param>
    /// <remarks>
    /// Creating a team vault is the only case: the client generates the key, wraps it to itself and
    /// sends the wrap, so the plaintext is already here and unwrapping the server's copy back would be
    /// a round trip to learn something this process just chose. Adopting it also means the new vault is
    /// usable immediately rather than at the next unlock, which is what somebody who just pressed
    /// "create" expects.
    /// </remarks>
    public void Adopt(Guid vaultId, byte[] vaultKey, uint keyGeneration)
    {
        ObjectDisposedException.ThrowIf(disposed, this);
        ArgumentNullException.ThrowIfNull(vaultKey);

        if (keys.TryGetValue(vaultId, out var previous))
        {
            CryptographicOperations.ZeroMemory(previous);
        }

        keys[vaultId] = vaultKey;
        generations[vaultId] = keyGeneration;

        Unopened = [.. Unopened.Where(id => id != vaultId)];
    }

    /// <summary>
    /// Tries to open a vault that has become readable since the session was unlocked.
    /// </summary>
    /// <returns>Whether the grant opened.</returns>
    /// <remarks>
    /// What a share looks like from the receiving end: the vault was in the list all along, listed and
    /// unreadable, and a member holding Share has now wrapped its key. Re-opening it here rather than
    /// waiting for a relock is the difference between "someone shared a vault with you" arriving and
    /// arriving tomorrow.
    /// </remarks>
    public bool TryAdmit(UserSecretBundle bundle, StoredVault vault)
    {
        ObjectDisposedException.ThrowIf(disposed, this);
        ArgumentNullException.ThrowIfNull(bundle);
        ArgumentNullException.ThrowIfNull(vault);

        if (vault.WrappedVaultKey is null)
        {
            return false;
        }

        if (keys.ContainsKey(vault.VaultId) && generations[vault.VaultId] == vault.KeyGeneration)
        {
            return true;
        }

        var key = VaultKeys.TryUnwrap(
            bundle.EncryptionKey, vault.WrappedVaultKey, vault.VaultId, vault.KeyGeneration);

        if (key is null)
        {
            return false;
        }

        Adopt(vault.VaultId, key, vault.KeyGeneration);

        return true;
    }

    /// <summary>Records that a vault cannot be read, so the interface can say so.</summary>
    /// <remarks>
    /// The counterpart of <see cref="TryAdmit"/> for the case where the grant did not open. Kept
    /// explicit rather than inferred from the absence of a key, because "no key" is also what a vault
    /// this session has never heard of looks like.
    /// </remarks>
    public void MarkUnreadable(Guid vaultId)
    {
        ObjectDisposedException.ThrowIf(disposed, this);

        if (!Unopened.Contains(vaultId))
        {
            Unopened = [.. Unopened, vaultId];
        }
    }

    /// <summary>
    /// Borrows a vault's key.
    /// </summary>
    /// <remarks>
    /// The returned memory is the keyring's own buffer, not a copy, and is zeroed when the keyring is
    /// disposed. Callers must not retain it past the operation they borrowed it for.
    /// </remarks>
    public bool TryGet(Guid vaultId, out ReadOnlyMemory<byte> vaultKey, out uint keyGeneration)
    {
        ObjectDisposedException.ThrowIf(disposed, this);

        if (keys.TryGetValue(vaultId, out var key))
        {
            vaultKey = key;
            keyGeneration = generations[vaultId];
            return true;
        }

        vaultKey = default;
        keyGeneration = 0;
        return false;
    }

    /// <summary>Whether this vault can be read at all.</summary>
    public bool CanRead(Guid vaultId) => !disposed && keys.ContainsKey(vaultId);

    /// <inheritdoc />
    public void Dispose()
    {
        if (disposed)
        {
            return;
        }

        disposed = true;

        foreach (var key in keys.Values)
        {
            CryptographicOperations.ZeroMemory(key);
        }

        keys.Clear();
        generations.Clear();
    }
}

/// <summary>Thrown when an operation needs a vault key the keyring does not hold.</summary>
/// <remarks>
/// An exception rather than a silent no-op, because every caller that reaches this point has already
/// been given the chance to check <see cref="VaultKeyring.CanRead"/>. Continuing without the key would
/// mean writing an item nobody can open.
/// </remarks>
[SuppressMessage(
    "Design",
    "CA1032:Implement standard exception constructors",
    Justification = "The vault id is required context; a message-only constructor would lose it.")]
public sealed class VaultUnreadableException(Guid vaultId)
    : InvalidOperationException(
        $"Vault {vaultId} has no usable key. Its grant is missing or awaiting re-wrap after a rekey.")
{
    /// <summary>The vault that cannot be read.</summary>
    public Guid VaultId { get; } = vaultId;
}