using System.Security.Cryptography;
using DodoSSH.Client.Auth;
using DodoSSH.Contracts;
using DodoSSH.Crypto;
using NSec.Cryptography;

namespace DodoSSH.Client.Api;

/// <summary>What enrolling produced, for the caller to hold and persist.</summary>
/// <remarks>
/// The bundle and the vault key are live secrets. The caller owns their lifetime and must dispose the
/// bundle; neither is ever written anywhere but the OS keystore and the encrypted local cache.
/// </remarks>
/// <param name="Response">The server's answer, including the vault and key log position.</param>
/// <param name="Bundle">The identity key pair, unlocked for this session.</param>
/// <param name="PersonalVaultKey">The personal vault's key, in plaintext for this session.</param>
/// <param name="DevicePrivateKey">
/// The enrolled device's X25519 private key, or <see langword="null"/> when no device was bound.
/// Belongs in the OS keystore — it is what lets a later launch unlock without the passphrase, and it
/// is the only thing that can open the device wrap held server-side.
/// </param>
/// <param name="RecoveryCode">
/// The generated recovery code, which must be shown to the user once and never stored. Losing this
/// along with the passphrase and every device means the vault is unrecoverable, and no server-side
/// reset is possible by design.
/// </param>
public sealed record EnrollmentOutcome(
    EnrollmentResponse Response,
    UserSecretBundle Bundle,
    byte[] PersonalVaultKey,
    byte[]? DevicePrivateKey,
    string RecoveryCode);

/// <summary>
/// Runs enrollment: generate keys, have the identity provider sign over them, and publish.
/// </summary>
/// <remarks>
/// <para>
/// Ordering here is not a matter of taste. The secret bundle's AAD binds to the server-assigned user
/// id, so <c>/me</c> must be read before anything can be wrapped — which is why <c>/me</c> provisions
/// the account and returns its id even when it reports that enrollment is required.
/// </para>
/// <para>
/// Everything the server receives is opaque to it. It gets public keys, wrapped blobs it cannot open,
/// and signatures it does not verify beyond the statement's own. That is the whole point: the server
/// stores the vault and cannot read it.
/// </para>
/// </remarks>
public sealed class ClientEnrollment(
    IAccountApi api,
    IKeyBindingAuthorizer keyBinding,
    TimeProvider clock,
    Argon2Profile? passphraseProfile = null)
{
    /// <remarks>
    /// Configurable because the cost is a product decision, not a constant: the plan exposes 128, 256 and
    /// 512 MiB security levels, and the parameters travel with the wrap so a user's choice is theirs
    /// alone. It also lets a test pay milliseconds instead of a third of a second to prove something that
    /// has nothing to do with how hard the passphrase is to attack.
    /// </remarks>
    private readonly Argon2Profile passphraseProfile = passphraseProfile ?? Argon2Profile.PassphraseDefault;

    /// <summary>Bytes of entropy behind a recovery code.</summary>
    private const int RecoveryEntropyBytes = 20;

    /// <summary>
    /// Enrolls the caller.
    /// </summary>
    /// <param name="me">
    /// The result of <see cref="DodoSshApiClient.GetMeAsync"/>, which supplies the user id the wraps
    /// bind to.
    /// </param>
    /// <param name="passphrase">The vault passphrase. Never transmitted or stored.</param>
    /// <param name="deviceName">Human-readable name for this machine.</param>
    /// <param name="vaultName">Display name for the personal vault. Plaintext, as vault names are.</param>
    /// <param name="bindThisDevice">
    /// Whether to register a device key so a later launch can unlock without the passphrase.
    /// <para>
    /// Pass <see langword="false"/> when the caller has nowhere durable to keep the private half. A
    /// device wrap whose private key does not survive the process is a row on the server that nobody can
    /// ever open, and it makes the account's device list claim a capability this machine does not have —
    /// which is worse than not offering it.
    /// </para>
    /// </param>
    /// <param name="cancellationToken">Cancellation token.</param>
    public async Task<EnrollmentOutcome> EnrollAsync(
        MeResponse me,
        string passphrase,
        string deviceName,
        string vaultName,
        bool bindThisDevice,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(me);
        ArgumentException.ThrowIfNullOrWhiteSpace(passphrase);
        ArgumentException.ThrowIfNullOrWhiteSpace(deviceName);

        var now = clock.GetUtcNow();
        var bundle = UserSecretBundle.Create(now);

        try
        {
            var statement = BuildStatement(me, bundle, now, deviceName);

            // The provider signs over the statement's hash, which is what stops the DodoSSH server
            // fabricating a key for a user who never enrolled. See ADR 0001.
            var idToken = await keyBinding
                .AuthorizeKeyBindingAsync(
                    KeyStatementCodec.ComputeNonce(ToFields(statement)),
                    cancellationToken)
                .ConfigureAwait(false);

            var request = BuildRequest(
                me, bundle, statement, idToken, passphrase, passphraseProfile, vaultName,
                bindThisDevice, now, out var material);

            var response = await api.EnrollAsync(request, cancellationToken).ConfigureAwait(false);

            return new EnrollmentOutcome(
                response,
                bundle,
                material.VaultKey,
                material.DevicePrivateKey,
                material.RecoveryCode);
        }
        catch
        {
            // The caller never receives the bundle on failure, so this is the only place that can
            // release its guarded memory.
            bundle.Dispose();
            throw;
        }
    }

    private static KeyStatement BuildStatement(
        MeResponse me,
        UserSecretBundle bundle,
        DateTimeOffset now,
        string deviceName) =>
        new(
            Version: KeyStatementCodec.CurrentVersion,
            Issuer: me.Issuer,
            Subject: me.Subject,
            Email: me.Email,
            EncryptionPublicKey: bundle.EncryptionPublicKey,
            SigningPublicKey: bundle.SigningPublicKey,
            KeyGeneration: 1,
            CreatedAt: now,
            DeviceName: deviceName);

    /// <summary>Secrets the caller keeps after a successful enrollment.</summary>
    private readonly record struct SessionMaterial(
        byte[] VaultKey,
        byte[]? DevicePrivateKey,
        string RecoveryCode);

    /// <remarks>
    /// The passphrase is a parameter rather than a field, so it lives only for the duration of this call
    /// and never becomes state on a long-lived object that a heap dump would find.
    /// </remarks>
    private static EnrollmentRequest BuildRequest(
        MeResponse me,
        UserSecretBundle bundle,
        KeyStatement statement,
        string idToken,
        string passphrase,
        Argon2Profile passphraseProfile,
        string vaultName,
        bool bindThisDevice,
        DateTimeOffset now,
        out SessionMaterial material)
    {
        var descriptor = DshAad.UserSecretBundle(me.UserId);

        // A fresh salt per wrap, and the parameters travel with it — so raising the cost later is a
        // per-user migration at next unlock rather than a breaking change.
        var passphraseSalt = RandomNumberGenerator.GetBytes(CryptoSpec.SaltSize);
        var recoverySalt = RandomNumberGenerator.GetBytes(CryptoSpec.SaltSize);
        var recoveryCode = GenerateRecoveryCode();

        byte[] passphraseWrap;
        byte[] recoveryWrap;

        using (var master = MasterKey.Derive(passphrase, passphraseSalt, passphraseProfile))
        {
            passphraseWrap = master.WrapBundle(bundle, descriptor);
        }

        // The recovery code carries real entropy, so it needs far less stretching than a passphrase.
        using (var recoveryMaster = MasterKey.Derive(
            recoveryCode, recoverySalt, Argon2Profile.RandomSecret))
        {
            recoveryWrap = recoveryMaster.WrapBundle(bundle, descriptor);
        }

        byte[]? devicePublicKey = null;
        byte[]? deviceWrap = null;
        byte[]? devicePrivateKey = null;

        if (bindThisDevice)
        {
            using var deviceKey = Key.Create(
                KeyAgreementAlgorithm.X25519,
                new KeyCreationParameters { ExportPolicy = KeyExportPolicies.AllowPlaintextExport });

            devicePublicKey = deviceKey.PublicKey.Export(KeyBlobFormat.RawPublicKey);
            deviceWrap = bundle.SealTo(devicePublicKey, descriptor);
            devicePrivateKey = deviceKey.Export(KeyBlobFormat.RawPrivateKey);
        }

        var vault = BuildPersonalVault(me, bundle, vaultName, now, out var vaultKey);

        material = new SessionMaterial(vaultKey, devicePrivateKey, recoveryCode);

        return new EnrollmentRequest(
            Statement: statement,
            StatementSignature: DshSignatures.SignKeyStatement(
                bundle.SigningKey,
                KeyStatementCodec.Encode(ToFields(statement))),
            IdentityProviderToken: idToken,
            WrappedPrivateKey: passphraseWrap,
            KdfParameters: ToContract(passphraseSalt, passphraseProfile),
            DevicePublicKey: devicePublicKey,
            DeviceWrappedPrivateKey: deviceWrap,
            RecoveryWrappedPrivateKey: recoveryWrap,
            RecoveryKdfParameters: ToContract(recoverySalt, Argon2Profile.RandomSecret),
            PersonalVault: vault);
    }

    /// <remarks>
    /// The vault id is chosen here rather than by the server, which is what makes enrollment safely
    /// retryable and is required by the grant signature — the signed tuple covers the vault id.
    /// <para>
    /// A self-grant carries no key log head: there is no third party whose key could have been
    /// substituted, and the log entry that would supply one is written by the server in the same
    /// transaction, so it cannot be signed over here.
    /// </para>
    /// </remarks>
    private static PersonalVaultRequest BuildPersonalVault(
        MeResponse me,
        UserSecretBundle bundle,
        string vaultName,
        DateTimeOffset now,
        out byte[] vaultKey)
    {
        var vaultId = Guid.CreateVersion7();
        vaultKey = VaultKeys.Create();

        var wrappedVaultKey = VaultKeys.WrapTo(vaultKey, bundle.EncryptionPublicKey, vaultId, 1);

        var fingerprint = DshCrypto.ComputeFingerprint(
            bundle.EncryptionPublicKey,
            bundle.SigningPublicKey);

        var grant = GrantStatementCodec.Encode(
            vaultId,
            keyGeneration: 1,
            GrantPurpose.Member,
            granteeUserId: me.UserId,
            granteeKeyFingerprint: fingerprint,
            wrappedKey: wrappedVaultKey,
            granterUserId: me.UserId,
            granterKeyFingerprint: fingerprint,
            keyLogHead: default,
            grantedAt: now);

        return new PersonalVaultRequest(
            VaultId: vaultId,
            Name: vaultName,
            WrappedVaultKey: wrappedVaultKey,
            GrantSignature: GrantStatementCodec.Sign(bundle.SigningKey, grant),
            GrantedAt: now);
    }

    /// <summary>
    /// Generates a printable recovery code.
    /// </summary>
    /// <remarks>
    /// Base32 over Crockford's alphabet, which omits I, L, O and U — so a code read aloud or copied off
    /// a screen cannot be mistranscribed into a different valid code, and cannot spell anything
    /// unfortunate. Grouped for legibility, and the groups are not part of the secret: the derivation
    /// uses the string exactly as shown, dashes included, because that is what the user will type back.
    /// </remarks>
    private static string GenerateRecoveryCode()
    {
        const string Alphabet = "0123456789ABCDEFGHJKMNPQRSTVWXYZ";
        const int BitsPerCharacter = 5;
        const int CharactersPerGroup = 5;

        // 20 bytes is 160 bits, which divides evenly into 32 five-bit characters — so no bits are
        // discarded and no padding is needed. Separators go between groups, hence one fewer than the
        // number of groups.
        var totalCharacters = RecoveryEntropyBytes * 8 / BitsPerCharacter;
        var separators = (totalCharacters - 1) / CharactersPerGroup;

        var entropy = RandomNumberGenerator.GetBytes(RecoveryEntropyBytes);
        var characters = new char[totalCharacters + separators];

        var index = 0;

        for (var position = 0; position < totalCharacters; position++)
        {
            if (position > 0 && position % CharactersPerGroup == 0)
            {
                characters[index++] = '-';
            }

            var value = 0;
            for (var offset = 0; offset < BitsPerCharacter; offset++)
            {
                var bit = (position * BitsPerCharacter) + offset;
                value = (value << 1) | ((entropy[bit / 8] >> (7 - (bit % 8))) & 1);
            }

            characters[index++] = Alphabet[value];
        }

        CryptographicOperations.ZeroMemory(entropy);

        return new string(characters);
    }

    private static KdfParameters ToContract(byte[] salt, Argon2Profile profile) =>
        new(
            Algorithm: "argon2id",
            Salt: salt,
            MemoryKibibytes: profile.MemoryKibibytes,
            Passes: profile.Passes,
            Parallelism: Argon2Profile.Parallelism);

    private static KeyStatementFields ToFields(KeyStatement statement) =>
        new(
            statement.Version,
            statement.Issuer,
            statement.Subject,
            statement.Email,
            statement.EncryptionPublicKey,
            statement.SigningPublicKey,
            statement.KeyGeneration,
            statement.CreatedAt,
            statement.DeviceName);
}