package vaultik_test

import (
	"bytes"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"io"
	"testing"

	"github.com/klauspost/compress/zstd"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"sneak.berlin/go/vaultik/internal/crypto"
)

// TestTeeReaderWithDecryption tests that TeeReader correctly hashes all encrypted
// bytes when streaming through age decryption and zstd decompression.
// This validates the verification path: hash encrypted blob -> decrypt -> decompress.
func TestTeeReaderWithDecryption(t *testing.T) {
	// Test data - use random data that doesn't compress well (5MB)
	testData := make([]byte, 5*1024*1024)
	_, err := rand.Read(testData)
	require.NoError(t, err)

	// Compress the data
	var compressedBuf bytes.Buffer

	compressor, err := zstd.NewWriter(&compressedBuf, zstd.WithEncoderLevel(zstd.SpeedDefault))
	require.NoError(t, err)
	_, err = compressor.Write(testData)
	require.NoError(t, err)
	err = compressor.Close()
	require.NoError(t, err)

	// Encrypt the compressed data
	testRecipient := "age1cplgrwj77ta54dnmydvvmzn64ltk83ankxl5sww04mrtmu62kv3s89gmvv"
	testSecretKey := "AGE-SECRET-KEY-1C77PYNTHXSHNNC6EYR2W52UWYXACXA5JT00J9CCW9986M3XY87PSGP89AQ"

	encryptor, err := crypto.NewEncryptor([]string{testRecipient})
	require.NoError(t, err)

	var encryptedBuf bytes.Buffer

	err = encryptor.EncryptStream(&encryptedBuf, bytes.NewReader(compressedBuf.Bytes()))
	require.NoError(t, err)

	encryptedData := encryptedBuf.Bytes()

	// Calculate the expected hash of the encrypted data directly
	expectedHash := sha256.Sum256(encryptedData)
	expectedHashStr := hex.EncodeToString(expectedHash[:])

	t.Logf("Encrypted data size: %d bytes", len(encryptedData))
	t.Logf("Expected hash: %s", expectedHashStr)

	// Now simulate what verifyBlob does: use TeeReader to hash while decrypting
	decryptor, err := crypto.NewDecryptor(testSecretKey)
	require.NoError(t, err)

	// Create hasher and tee reader
	hasher := sha256.New()
	reader := bytes.NewReader(encryptedData)
	teeReader := io.TeeReader(reader, hasher)

	// Decrypt through the tee reader
	decryptedReader, err := decryptor.DecryptStream(teeReader)
	require.NoError(t, err)

	// Decompress
	decompressor, err := zstd.NewReader(decryptedReader)
	require.NoError(t, err)

	defer decompressor.Close()

	// Read all decompressed data (simulating chunk verification)
	decompressedData, err := io.ReadAll(decompressor)
	require.NoError(t, err)

	// Verify we got the original data back
	assert.Equal(t, testData, decompressedData, "Decompressed data should match original")

	// Drain remaining decompressed data (should be 0)
	remaining, err := io.Copy(io.Discard, decompressor)
	require.NoError(t, err)
	assert.Equal(t, int64(0), remaining, "No remaining decompressed data")

	// Calculate hash from tee reader
	calculatedHashStr := hex.EncodeToString(hasher.Sum(nil))
	t.Logf("Calculated hash (before drain): %s", calculatedHashStr)

	// Verify the hash matches the direct hash of encrypted data
	assert.Equal(t, expectedHashStr, calculatedHashStr,
		"Hash calculated via TeeReader should match direct hash of encrypted data")
}