package blobgen_test

import (
	"bytes"
	"crypto/sha256"
	"fmt"
	"io"
	"testing"

	"filippo.io/age"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"sneak.berlin/go/vaultik/internal/blobgen"
)

// checkRoundTrip writes input through a Writer, reads it back through a Reader,
// and verifies the plaintext, the byte counts, and the content hashes.
func checkRoundTrip(
	t *testing.T, id *age.X25519Identity, recipient string,
	level int, input []byte,
) {
	t.Helper()

	var buf bytes.Buffer

	w, err := blobgen.NewWriter(&buf, level, []string{recipient})
	require.NoError(t, err)

	n, err := w.Write(input)
	require.NoError(t, err)
	assert.Equal(t, len(input), n)
	require.NoError(t, w.Close())
	require.Equal(t, int64(len(input)), w.BytesWritten())

	r, err := blobgen.NewReader(bytes.NewReader(buf.Bytes()), id)
	require.NoError(t, err)

	got, err := io.ReadAll(r)
	require.NoError(t, err)
	require.NoError(t, r.Close())

	assert.Equal(t, input, got, "decrypted output must equal input")
	require.Equal(t, int64(len(input)), r.BytesRead())

	// The hash values are checked by decrypting: the reader's single SHA-256
	// is the hash of the plaintext, and hashing it once more (DoubleSHA256)
	// gives the writer's ContentID.
	single := sha256.Sum256(got)
	assert.Equal(t, single[:], r.Sum256())
	assert.Equal(t, blobgen.DoubleSHA256(r.Sum256()), w.ContentID())
}

// TestWriterReaderRoundTrip covers issue cases 1 and 2: every size round trips
// for both random and compressible data, and the reader hash, its double hash
// and the byte counts all agree.
func TestWriterReaderRoundTrip(t *testing.T) {
	t.Parallel()

	id, recipient := makeIdentity(t)

	// Sizes exercise the age segment boundary (64 KiB) from just below to a
	// few segments above it, plus the empty and single-byte edges.
	sizes := []int{0, 1, 65535, 65536, 65537, 4*65536 + 123}

	kinds := []struct {
		name string
		fill func(*testing.T, int) []byte
	}{
		{"random", randomBytes},
		{"compressible", func(_ *testing.T, n int) []byte {
			return compressibleBytes(n)
		}},
	}

	for _, k := range kinds {
		for _, size := range sizes {
			name := fmt.Sprintf("%s/%d", k.name, size)
			t.Run(name, func(t *testing.T) {
				t.Parallel()
				checkRoundTrip(t, id, recipient, 1, k.fill(t, size))
			})
		}
	}
}

// TestZeroLengthNoWrite covers issue case 3: a Writer closed with no Write at
// all produces the double hash of the empty input, and the blob reads back as
// empty with no error.
func TestZeroLengthNoWrite(t *testing.T) {
	t.Parallel()

	id, recipient := makeIdentity(t)

	var buf bytes.Buffer

	w, err := blobgen.NewWriter(&buf, 1, []string{recipient})
	require.NoError(t, err)
	require.NoError(t, w.Close())
	assert.Equal(t, int64(0), w.BytesWritten())

	empty := sha256.Sum256(nil)
	doubled := sha256.Sum256(empty[:])
	assert.Equal(t, doubled[:], w.ContentID(),
		"ContentID of empty input is SHA256(SHA256(\"\"))")

	r, err := blobgen.NewReader(bytes.NewReader(buf.Bytes()), id)
	require.NoError(t, err)

	got, err := io.ReadAll(r)
	require.NoError(t, err)
	require.NoError(t, r.Close())

	assert.Empty(t, got, "empty blob decrypts to empty output")
	assert.Equal(t, int64(0), r.BytesRead())
	assert.Equal(t, empty[:], r.Sum256())
}

// TestNewWriterValidLevelsRoundTrip covers the accepted end of issue case 9:
// the boundary compression levels 1 and 19 both round trip.
func TestNewWriterValidLevelsRoundTrip(t *testing.T) {
	t.Parallel()

	id, recipient := makeIdentity(t)
	input := randomBytes(t, 4096)

	for _, level := range []int{1, 19} {
		t.Run(fmt.Sprintf("level%d", level), func(t *testing.T) {
			t.Parallel()
			checkRoundTrip(t, id, recipient, level, input)
		})
	}
}