using EntKube.Secrets.Crypto;
using FluentAssertions;

namespace EntKube.Secrets.Tests.Crypto;

/// <summary>
/// Tests for Shamir's Secret Sharing scheme. This is the mechanism that protects
/// the master encryption key — the key is split into N shares, and any M of those
/// shares can reconstruct it. Fewer than M shares reveal nothing about the key.
///
/// This is critical for the vault's seal/unseal lifecycle: on initialization the
/// master key is split into shares distributed to key holders. To unseal the vault
/// after a restart, M key holders must each provide their share.
/// </summary>
public class ShamirSecretSharingTests
{
    [Fact]
    public void Split_ThenCombine_WithExactThreshold_RecoversSecret()
    {
        // Arrange — A 256-bit master key, split into 5 shares with a threshold of 3.
        // Any 3 shares should be enough to recover the original key.

        byte[] secret = AesGcmEncryptor.GenerateKey();

        // Act — Split into 5 shares, then recombine using exactly 3.

        List<ShamirShare> shares = ShamirSecretSharing.Split(secret, totalShares: 5, threshold: 3);
        byte[] recovered = ShamirSecretSharing.Combine(shares.Take(3).ToList(), threshold: 3);

        // Assert — The recovered secret must match the original exactly.

        recovered.Should().BeEquivalentTo(secret);
    }

    [Fact]
    public void Split_ThenCombine_WithMoreThanThreshold_RecoversSecret()
    {
        // Arrange — Using more shares than required should also work.
        // This tests that extra shares don't corrupt the reconstruction.

        byte[] secret = AesGcmEncryptor.GenerateKey();

        // Act — Split into 5, recombine using all 5 (threshold is still 3).

        List<ShamirShare> shares = ShamirSecretSharing.Split(secret, totalShares: 5, threshold: 3);
        byte[] recovered = ShamirSecretSharing.Combine(shares, threshold: 3);

        // Assert

        recovered.Should().BeEquivalentTo(secret);
    }

    [Fact]
    public void Split_ThenCombine_WithDifferentShareSubsets_AllRecoverSecret()
    {
        // Arrange — Any combination of M-of-N shares should work, not just
        // the first M. This tests that shares 2,3,5 work as well as 1,2,3.

        byte[] secret = AesGcmEncryptor.GenerateKey();
        List<ShamirShare> shares = ShamirSecretSharing.Split(secret, totalShares: 5, threshold: 3);

        // Act & Assert — Multiple subsets of 3 shares all recover the secret.

        ShamirSecretSharing.Combine(new List<ShamirShare> { shares[0], shares[1], shares[2] }, 3)
            .Should().BeEquivalentTo(secret, "shares 1,2,3 should recover the secret");

        ShamirSecretSharing.Combine(new List<ShamirShare> { shares[0], shares[2], shares[4] }, 3)
            .Should().BeEquivalentTo(secret, "shares 1,3,5 should recover the secret");

        ShamirSecretSharing.Combine(new List<ShamirShare> { shares[1], shares[3], shares[4] }, 3)
            .Should().BeEquivalentTo(secret, "shares 2,4,5 should recover the secret");
    }

    [Fact]
    public void Split_ProducesCorrectNumberOfShares()
    {
        // Arrange

        byte[] secret = AesGcmEncryptor.GenerateKey();

        // Act

        List<ShamirShare> shares = ShamirSecretSharing.Split(secret, totalShares: 7, threshold: 4);

        // Assert — Must produce exactly the requested number of shares.

        shares.Should().HaveCount(7);
    }

    [Fact]
    public void Split_EachShareHasUniqueIndex()
    {
        // Arrange

        byte[] secret = AesGcmEncryptor.GenerateKey();

        // Act

        List<ShamirShare> shares = ShamirSecretSharing.Split(secret, totalShares: 5, threshold: 3);

        // Assert — Share indices must be unique (used as x-coordinates in Lagrange interpolation).

        List<int> indices = shares.Select(s => s.Index).ToList();
        indices.Should().OnlyHaveUniqueItems();
    }

    [Fact]
    public void Split_WithThresholdOf1_EachShareIsTheSecret()
    {
        // Arrange — Threshold of 1 means any single share recovers the secret.
        // This is the degenerate case (no split protection).

        byte[] secret = AesGcmEncryptor.GenerateKey();

        // Act

        List<ShamirShare> shares = ShamirSecretSharing.Split(secret, totalShares: 3, threshold: 1);

        // Assert — Each individual share should recover the secret.

        foreach (ShamirShare share in shares)
        {
            byte[] recovered = ShamirSecretSharing.Combine(new List<ShamirShare> { share }, threshold: 1);
            recovered.Should().BeEquivalentTo(secret);
        }
    }

    [Fact]
    public void Split_WithThresholdEqualsTotal_RequiresAllShares()
    {
        // Arrange — 3-of-3 means all shares are needed.

        byte[] secret = AesGcmEncryptor.GenerateKey();

        // Act

        List<ShamirShare> shares = ShamirSecretSharing.Split(secret, totalShares: 3, threshold: 3);
        byte[] recovered = ShamirSecretSharing.Combine(shares, threshold: 3);

        // Assert

        recovered.Should().BeEquivalentTo(secret);
    }

    [Fact]
    public void Split_WithInvalidParameters_Throws()
    {
        // Threshold must be >= 1, total must be >= threshold.

        byte[] secret = AesGcmEncryptor.GenerateKey();

        Action zeroThreshold = () => ShamirSecretSharing.Split(secret, totalShares: 5, threshold: 0);
        zeroThreshold.Should().Throw<ArgumentException>();

        Action thresholdExceedsTotal = () => ShamirSecretSharing.Split(secret, totalShares: 3, threshold: 5);
        thresholdExceedsTotal.Should().Throw<ArgumentException>();
    }

    [Fact]
    public void Combine_WithTooFewShares_ProducesWrongResult()
    {
        // Arrange — With fewer than threshold shares, Lagrange interpolation
        // produces a different polynomial, so the result should not match.
        // This is the information-theoretic security guarantee of Shamir's scheme.

        byte[] secret = AesGcmEncryptor.GenerateKey();
        List<ShamirShare> shares = ShamirSecretSharing.Split(secret, totalShares: 5, threshold: 3);

        // Act — Try to combine only 2 shares when 3 are needed.

        byte[] wrongResult = ShamirSecretSharing.Combine(shares.Take(2).ToList(), threshold: 2);

        // Assert — The result must NOT equal the original secret.

        wrongResult.Should().NotBeEquivalentTo(secret);
    }

    [Fact]
    public void Split_ThenCombine_WorksWithDifferentSecretSizes()
    {
        // Shamir's scheme works byte-by-byte, so it should handle any size secret.

        byte[] smallSecret = new byte[] { 0x42 };
        byte[] largeSecret = new byte[64]; // 512-bit key
        Random.Shared.NextBytes(largeSecret);

        List<ShamirShare> smallShares = ShamirSecretSharing.Split(smallSecret, 3, 2);
        List<ShamirShare> largeShares = ShamirSecretSharing.Split(largeSecret, 3, 2);

        ShamirSecretSharing.Combine(smallShares.Take(2).ToList(), 2)
            .Should().BeEquivalentTo(smallSecret);

        ShamirSecretSharing.Combine(largeShares.Take(2).ToList(), 2)
            .Should().BeEquivalentTo(largeSecret);
    }
}