// Copyright © StellaOps. All rights reserved.
// SPDX-License-Identifier: BUSL-1.1
// Sprint: SPRINT_20260112_018_CRYPTO_key_escrow_shamir
// Tasks: ESCROW-004, ESCROW-006, ESCROW-008, ESCROW-009


using System.Security.Cryptography;

namespace StellaOps.Cryptography.KeyEscrow;

/// <summary>
/// Implementation of key escrow service using Shamir's Secret Sharing.
/// </summary>
public sealed class KeyEscrowService : IKeyEscrowService
{
    private readonly IEscrowAgentStore _agentStore;
    private readonly IKeyEscrowAuditLogger _auditLogger;
    private readonly ShamirSecretSharing _shamir;
    private readonly TimeProvider _timeProvider;
    private readonly KeyEscrowServiceOptions _options;

    public KeyEscrowService(
        IEscrowAgentStore agentStore,
        IKeyEscrowAuditLogger auditLogger,
        TimeProvider timeProvider,
        KeyEscrowServiceOptions? options = null)
    {
        _agentStore = agentStore ?? throw new ArgumentNullException(nameof(agentStore));
        _auditLogger = auditLogger ?? throw new ArgumentNullException(nameof(auditLogger));
        _timeProvider = timeProvider ?? throw new ArgumentNullException(nameof(timeProvider));
        _options = options ?? new KeyEscrowServiceOptions();
        _shamir = new ShamirSecretSharing();
    }

    /// <inheritdoc />
    public async Task<KeyEscrowResult> EscrowKeyAsync(
        string keyId,
        byte[] keyMaterial,
        KeyEscrowOptions options,
        CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(keyId);
        ArgumentNullException.ThrowIfNull(keyMaterial);
        ArgumentNullException.ThrowIfNull(options);

        var now = _timeProvider.GetUtcNow();
        var expiresAt = now.AddDays(options.ExpirationDays);

        try
        {
            // Get agents to distribute shares to
            var agents = await GetAgentsForDistributionAsync(options, cancellationToken);
            if (agents.Count < options.TotalShares)
            {
                return CreateFailureResult(keyId, $"Insufficient agents: need {options.TotalShares}, have {agents.Count}");
            }

            // Split the key
            var shamirShares = _shamir.Split(keyMaterial, options.Threshold, options.TotalShares);

            // Create and store encrypted shares
            var shareIds = new List<Guid>();
            var custodianIds = new List<string>();

            for (int i = 0; i < shamirShares.Length; i++)
            {
                var agent = agents[i];
                var shamirShare = shamirShares[i];

                // Encrypt share for agent
                var (encryptedData, encryptionInfo) = await EncryptShareAsync(
                    shamirShare.Data,
                    agent,
                    cancellationToken);

                // Compute checksum of unencrypted data
                var checksum = ComputeChecksum(shamirShare.Data);

                var keyShare = new KeyShare
                {
                    ShareId = Guid.NewGuid(),
                    Index = shamirShare.Index,
                    EncryptedData = encryptedData,
                    KeyId = keyId,
                    Threshold = options.Threshold,
                    TotalShares = options.TotalShares,
                    CreatedAt = now,
                    ExpiresAt = expiresAt,
                    CustodianId = agent.AgentId,
                    ChecksumHex = checksum,
                    EncryptionInfo = encryptionInfo,
                };

                await _agentStore.StoreShareAsync(keyShare, cancellationToken);
                shareIds.Add(keyShare.ShareId);
                custodianIds.Add(agent.AgentId);

                // Clear sensitive data
                Array.Clear(shamirShare.Data);
            }

            // Store metadata
            var metadata = new KeyEscrowMetadata
            {
                KeyId = keyId,
                Threshold = options.Threshold,
                TotalShares = options.TotalShares,
                CreatedAt = now,
                ExpiresAt = expiresAt,
                RequireDualControl = options.RequireDualControl,
                CustodianIds = custodianIds,
                Metadata = options.Metadata,
            };

            await _agentStore.StoreEscrowMetadataAsync(metadata, cancellationToken);

            // Audit log
            await _auditLogger.LogEscrowAsync(new KeyEscrowAuditEvent
            {
                EventId = Guid.NewGuid(),
                EventType = KeyEscrowAuditEventType.KeyEscrowed,
                KeyId = keyId,
                Timestamp = now,
                InitiatorId = "system", // TODO: get from context
                CustodianIds = custodianIds,
                ShareCount = options.TotalShares,
                Success = true,
            }, cancellationToken);

            return new KeyEscrowResult
            {
                Success = true,
                KeyId = keyId,
                ShareIds = shareIds,
                Threshold = options.Threshold,
                TotalShares = options.TotalShares,
                ExpiresAt = expiresAt,
            };
        }
        catch (Exception ex)
        {
            await _auditLogger.LogEscrowAsync(new KeyEscrowAuditEvent
            {
                EventId = Guid.NewGuid(),
                EventType = KeyEscrowAuditEventType.KeyEscrowed,
                KeyId = keyId,
                Timestamp = now,
                InitiatorId = "system",
                Success = false,
                Error = ex.Message,
            }, cancellationToken);

            return CreateFailureResult(keyId, ex.Message);
        }
    }

    /// <inheritdoc />
    public async Task<KeyRecoveryResult> RecoverKeyAsync(
        KeyRecoveryRequest request,
        IReadOnlyList<KeyShare> shares,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(request);
        ArgumentNullException.ThrowIfNull(shares);

        var now = _timeProvider.GetUtcNow();

        try
        {
            // Get escrow metadata
            var metadata = await _agentStore.GetEscrowMetadataAsync(request.KeyId, cancellationToken);
            if (metadata == null)
            {
                return CreateRecoveryFailure(request.KeyId, "Key not found in escrow");
            }

            // Validate share count
            if (shares.Count < metadata.Threshold)
            {
                return CreateRecoveryFailure(
                    request.KeyId,
                    $"Insufficient shares: need {metadata.Threshold}, have {shares.Count}");
            }

            // Validate authorizing custodians
            if (metadata.RequireDualControl && request.AuthorizingCustodians.Count < 2)
            {
                return CreateRecoveryFailure(
                    request.KeyId,
                    "Dual-control required: at least 2 custodians must authorize");
            }

            // Decrypt and verify shares
            var shamirShares = new List<ShamirShare>();
            foreach (var share in shares.Take(metadata.Threshold))
            {
                // In production, shares would be decrypted here
                // For now, assume EncryptedData contains decrypted share data (test scenario)
                var decryptedData = share.EncryptedData; // TODO: decrypt based on EncryptionInfo

                // Verify checksum
                var checksum = ComputeChecksum(decryptedData);
                if (checksum != share.ChecksumHex)
                {
                    return CreateRecoveryFailure(request.KeyId, $"Share {share.Index} failed checksum verification");
                }

                shamirShares.Add(new ShamirShare
                {
                    Index = (byte)share.Index,
                    Data = decryptedData,
                });
            }

            // Reconstruct the key
            var keyMaterial = _shamir.Combine(shamirShares.ToArray());

            var auditEventId = Guid.NewGuid();

            // Audit log
            await _auditLogger.LogEscrowAsync(new KeyEscrowAuditEvent
            {
                EventId = auditEventId,
                EventType = KeyEscrowAuditEventType.KeyRecovered,
                KeyId = request.KeyId,
                Timestamp = now,
                InitiatorId = request.InitiatorId,
                Reason = request.Reason,
                CustodianIds = request.AuthorizingCustodians.ToList(),
                ShareCount = shares.Count,
                Success = true,
                CeremonyId = request.CeremonyId,
            }, cancellationToken);

            return new KeyRecoveryResult
            {
                Success = true,
                KeyId = request.KeyId,
                KeyMaterial = keyMaterial,
                SharesUsed = shamirShares.Count,
                AuditEventId = auditEventId,
            };
        }
        catch (Exception ex)
        {
            await _auditLogger.LogEscrowAsync(new KeyEscrowAuditEvent
            {
                EventId = Guid.NewGuid(),
                EventType = KeyEscrowAuditEventType.RecoveryFailed,
                KeyId = request.KeyId,
                Timestamp = now,
                InitiatorId = request.InitiatorId,
                Reason = request.Reason,
                Success = false,
                Error = ex.Message,
            }, cancellationToken);

            return CreateRecoveryFailure(request.KeyId, ex.Message);
        }
    }

    /// <inheritdoc />
    public async Task<KeyEscrowStatus?> GetEscrowStatusAsync(
        string keyId,
        CancellationToken cancellationToken = default)
    {
        var metadata = await _agentStore.GetEscrowMetadataAsync(keyId, cancellationToken);
        if (metadata == null)
        {
            return null;
        }

        var shares = await _agentStore.GetSharesForKeyAsync(keyId, cancellationToken);
        var now = _timeProvider.GetUtcNow();
        var validShares = shares.Count(s => s.ExpiresAt > now);

        return new KeyEscrowStatus
        {
            KeyId = keyId,
            IsEscrowed = validShares > 0,
            Threshold = metadata.Threshold,
            TotalShares = metadata.TotalShares,
            ValidShares = validShares,
            CreatedAt = metadata.CreatedAt,
            ExpiresAt = metadata.ExpiresAt,
            CustodianIds = metadata.CustodianIds.ToList(),
        };
    }

    /// <inheritdoc />
    public async Task<IReadOnlyList<KeyEscrowSummary>> ListEscrowedKeysAsync(
        CancellationToken cancellationToken = default)
    {
        var keyIds = await _agentStore.ListEscrowedKeyIdsAsync(cancellationToken);
        var summaries = new List<KeyEscrowSummary>();

        foreach (var keyId in keyIds)
        {
            var metadata = await _agentStore.GetEscrowMetadataAsync(keyId, cancellationToken);
            if (metadata != null)
            {
                summaries.Add(new KeyEscrowSummary
                {
                    KeyId = keyId,
                    Threshold = metadata.Threshold,
                    TotalShares = metadata.TotalShares,
                    CreatedAt = metadata.CreatedAt,
                    ExpiresAt = metadata.ExpiresAt,
                    Metadata = metadata.Metadata,
                });
            }
        }

        return summaries;
    }

    /// <inheritdoc />
    public async Task<bool> RevokeEscrowAsync(
        string keyId,
        string reason,
        CancellationToken cancellationToken = default)
    {
        var now = _timeProvider.GetUtcNow();

        var deleted = await _agentStore.DeleteSharesForKeyAsync(keyId, cancellationToken);

        await _auditLogger.LogEscrowAsync(new KeyEscrowAuditEvent
        {
            EventId = Guid.NewGuid(),
            EventType = KeyEscrowAuditEventType.EscrowRevoked,
            KeyId = keyId,
            Timestamp = now,
            InitiatorId = "system", // TODO: get from context
            Reason = reason,
            ShareCount = deleted,
            Success = deleted > 0,
        }, cancellationToken);

        return deleted > 0;
    }

    /// <inheritdoc />
    public async Task<KeyEscrowResult> ReEscrowKeyAsync(
        string keyId,
        byte[] keyMaterial,
        KeyEscrowOptions? options = null,
        CancellationToken cancellationToken = default)
    {
        // Get existing metadata if no options provided
        if (options == null)
        {
            var existing = await _agentStore.GetEscrowMetadataAsync(keyId, cancellationToken);
            if (existing == null)
            {
                return CreateFailureResult(keyId, "No existing escrow found and no options provided");
            }

            options = new KeyEscrowOptions
            {
                Threshold = existing.Threshold,
                TotalShares = existing.TotalShares,
                RequireDualControl = existing.RequireDualControl,
                Metadata = existing.Metadata,
            };
        }

        // Revoke existing shares
        await _agentStore.DeleteSharesForKeyAsync(keyId, cancellationToken);

        // Create new escrow
        var result = await EscrowKeyAsync(keyId, keyMaterial, options, cancellationToken);

        if (result.Success)
        {
            await _auditLogger.LogEscrowAsync(new KeyEscrowAuditEvent
            {
                EventId = Guid.NewGuid(),
                EventType = KeyEscrowAuditEventType.KeyReEscrowed,
                KeyId = keyId,
                Timestamp = _timeProvider.GetUtcNow(),
                InitiatorId = "system",
                ShareCount = result.TotalShares,
                Success = true,
            }, cancellationToken);
        }

        return result;
    }

    private async Task<IReadOnlyList<EscrowAgent>> GetAgentsForDistributionAsync(
        KeyEscrowOptions options,
        CancellationToken cancellationToken)
    {
        if (options.AgentIds != null && options.AgentIds.Count >= options.TotalShares)
        {
            var agents = new List<EscrowAgent>();
            foreach (var agentId in options.AgentIds.Take(options.TotalShares))
            {
                var agent = await _agentStore.GetAgentAsync(agentId, cancellationToken);
                if (agent != null && agent.IsActive)
                {
                    agents.Add(agent);
                }
            }

            return agents;
        }

        return await _agentStore.GetActiveAgentsAsync(cancellationToken);
    }

    private Task<(byte[] EncryptedData, ShareEncryptionInfo Info)> EncryptShareAsync(
        byte[] shareData,
        EscrowAgent agent,
        CancellationToken cancellationToken)
    {
        // For now, use AES-256-GCM with a randomly generated key
        // In production, this would encrypt with the agent's public key

        using var aes = new AesGcm(GenerateKey(), AesGcm.TagByteSizes.MaxSize);
        var nonce = new byte[AesGcm.NonceByteSizes.MaxSize];
        RandomNumberGenerator.Fill(nonce);

        var ciphertext = new byte[shareData.Length];
        var tag = new byte[AesGcm.TagByteSizes.MaxSize];

        aes.Encrypt(nonce, shareData, ciphertext, tag);

        // Combine ciphertext and tag
        var encryptedData = new byte[ciphertext.Length + tag.Length];
        Buffer.BlockCopy(ciphertext, 0, encryptedData, 0, ciphertext.Length);
        Buffer.BlockCopy(tag, 0, encryptedData, ciphertext.Length, tag.Length);

        var info = new ShareEncryptionInfo
        {
            Algorithm = "AES-256-GCM",
            NonceBase64 = Convert.ToBase64String(nonce),
            AuthTagBase64 = Convert.ToBase64String(tag),
        };

        return Task.FromResult((encryptedData, info));
    }

    private static byte[] GenerateKey()
    {
        var key = new byte[32]; // 256 bits
        RandomNumberGenerator.Fill(key);
        return key;
    }

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

    private static KeyEscrowResult CreateFailureResult(string keyId, string error)
    {
        return new KeyEscrowResult
        {
            Success = false,
            KeyId = keyId,
            ShareIds = Array.Empty<Guid>(),
            Threshold = 0,
            TotalShares = 0,
            ExpiresAt = DateTimeOffset.MinValue,
            Error = error,
        };
    }

    private static KeyRecoveryResult CreateRecoveryFailure(string keyId, string error)
    {
        return new KeyRecoveryResult
        {
            Success = false,
            KeyId = keyId,
            Error = error,
        };
    }
}

/// <summary>
/// Options for the key escrow service.
/// </summary>
public sealed record KeyEscrowServiceOptions
{
    /// <summary>
    /// Default threshold for M-of-N splitting.
    /// </summary>
    public int DefaultThreshold { get; init; } = 3;

    /// <summary>
    /// Default total shares for M-of-N splitting.
    /// </summary>
    public int DefaultTotalShares { get; init; } = 5;

    /// <summary>
    /// Default expiration in days.
    /// </summary>
    public int DefaultExpirationDays { get; init; } = 365;

    /// <summary>
    /// Whether to automatically delete shares after recovery.
    /// </summary>
    public bool AutoDeleteOnRecovery { get; init; } = false;
}