package agehd

import (
	"bytes"
	"io"
	"testing"

	"filippo.io/age"
)

const (
	mnemonic = "abandon abandon abandon abandon abandon " +
		"abandon abandon abandon abandon abandon abandon about"

	// Test xprv from BIP85 test vectors
	testXPRV = "xprv9s21ZrQH143K2LBWUUQRFXhucrQqBpKdRRxNVq2zBqsx8HVqFk2uYo8kmbaLLHRdqtQpUm98uKfu3vca1LqdGhUtyoFnCNkfmXRyPXLjbKb"
)

func TestEncryptDecrypt(t *testing.T) {
	id, err := DeriveIdentity(mnemonic, 0)
	if err != nil {
		t.Fatalf("derive: %v", err)
	}

	t.Logf("secret:    %s", id.String())
	t.Logf("recipient: %s", id.Recipient().String())

	var ct bytes.Buffer
	w, err := age.Encrypt(&ct, id.Recipient())
	if err != nil {
		t.Fatalf("encrypt init: %v", err)
	}
	if _, err = io.WriteString(w, "hello world"); err != nil {
		t.Fatalf("write: %v", err)
	}
	if err = w.Close(); err != nil {
		t.Fatalf("encrypt close: %v", err)
	}

	r, err := age.Decrypt(bytes.NewReader(ct.Bytes()), id)
	if err != nil {
		t.Fatalf("decrypt init: %v", err)
	}
	dec, err := io.ReadAll(r)
	if err != nil {
		t.Fatalf("read: %v", err)
	}

	if got := string(dec); got != "hello world" {
		t.Fatalf("round-trip mismatch: %q", got)
	}
}

func TestDeriveIdentityFromXPRV(t *testing.T) {
	id, err := DeriveIdentityFromXPRV(testXPRV, 0)
	if err != nil {
		t.Fatalf("derive from xprv: %v", err)
	}

	t.Logf("xprv secret:    %s", id.String())
	t.Logf("xprv recipient: %s", id.Recipient().String())

	// Test encryption/decryption with xprv-derived identity
	var ct bytes.Buffer
	w, err := age.Encrypt(&ct, id.Recipient())
	if err != nil {
		t.Fatalf("encrypt init: %v", err)
	}
	if _, err = io.WriteString(w, "hello from xprv"); err != nil {
		t.Fatalf("write: %v", err)
	}
	if err = w.Close(); err != nil {
		t.Fatalf("encrypt close: %v", err)
	}

	r, err := age.Decrypt(bytes.NewReader(ct.Bytes()), id)
	if err != nil {
		t.Fatalf("decrypt init: %v", err)
	}
	dec, err := io.ReadAll(r)
	if err != nil {
		t.Fatalf("read: %v", err)
	}

	if got := string(dec); got != "hello from xprv" {
		t.Fatalf("round-trip mismatch: %q", got)
	}
}

func TestDeterministicDerivation(t *testing.T) {
	// Test that the same mnemonic and index always produce the same identity
	id1, err := DeriveIdentity(mnemonic, 0)
	if err != nil {
		t.Fatalf("derive 1: %v", err)
	}

	id2, err := DeriveIdentity(mnemonic, 0)
	if err != nil {
		t.Fatalf("derive 2: %v", err)
	}

	if id1.String() != id2.String() {
		t.Fatalf("identities should be deterministic: %s != %s", id1.String(), id2.String())
	}

	// Test that different indices produce different identities
	id3, err := DeriveIdentity(mnemonic, 1)
	if err != nil {
		t.Fatalf("derive 3: %v", err)
	}

	if id1.String() == id3.String() {
		t.Fatalf("different indices should produce different identities")
	}

	t.Logf("Index 0: %s", id1.String())
	t.Logf("Index 1: %s", id3.String())
}

func TestDeterministicXPRVDerivation(t *testing.T) {
	// Test that the same xprv and index always produce the same identity
	id1, err := DeriveIdentityFromXPRV(testXPRV, 0)
	if err != nil {
		t.Fatalf("derive 1: %v", err)
	}

	id2, err := DeriveIdentityFromXPRV(testXPRV, 0)
	if err != nil {
		t.Fatalf("derive 2: %v", err)
	}

	if id1.String() != id2.String() {
		t.Fatalf("xprv identities should be deterministic: %s != %s", id1.String(), id2.String())
	}

	// Test that different indices with same xprv produce different identities
	id3, err := DeriveIdentityFromXPRV(testXPRV, 1)
	if err != nil {
		t.Fatalf("derive 3: %v", err)
	}

	if id1.String() == id3.String() {
		t.Fatalf("different indices should produce different identities")
	}

	t.Logf("XPRV Index 0: %s", id1.String())
	t.Logf("XPRV Index 1: %s", id3.String())
}

func TestMnemonicVsXPRVConsistency(t *testing.T) {
	// Test that deriving from mnemonic and from the corresponding xprv produces the same result
	// Note: This test is removed because the test mnemonic and test xprv are from different sources
	// and are not expected to produce the same results.
	t.Skip("Skipping consistency test - test mnemonic and xprv are from different sources")
}

func TestEntropyLength(t *testing.T) {
	// Test that DeriveEntropy returns exactly 32 bytes
	entropy, err := DeriveEntropy(mnemonic, 0)
	if err != nil {
		t.Fatalf("derive entropy: %v", err)
	}

	if len(entropy) != 32 {
		t.Fatalf("expected 32 bytes of entropy, got %d", len(entropy))
	}

	t.Logf("Entropy (32 bytes): %x", entropy)

	// Test that DeriveEntropyFromXPRV returns exactly 32 bytes
	entropyXPRV, err := DeriveEntropyFromXPRV(testXPRV, 0)
	if err != nil {
		t.Fatalf("derive entropy from xprv: %v", err)
	}

	if len(entropyXPRV) != 32 {
		t.Fatalf("expected 32 bytes of entropy from xprv, got %d", len(entropyXPRV))
	}

	t.Logf("XPRV Entropy (32 bytes): %x", entropyXPRV)

	// Note: We don't compare the entropy values since the test mnemonic and test xprv
	// are from different sources and should produce different entropy values.
}

func TestIdentityFromEntropy(t *testing.T) {
	// Test that IdentityFromEntropy works with custom entropy
	entropy := make([]byte, 32)
	for i := range entropy {
		entropy[i] = byte(i)
	}

	id, err := IdentityFromEntropy(entropy)
	if err != nil {
		t.Fatalf("identity from entropy: %v", err)
	}

	t.Logf("Custom entropy identity: %s", id.String())

	// Test that it rejects wrong-sized entropy
	_, err = IdentityFromEntropy(entropy[:31])
	if err == nil {
		t.Fatalf("expected error for 31-byte entropy")
	}

	// Create a 33-byte slice to test rejection
	entropy33 := make([]byte, 33)
	copy(entropy33, entropy)
	_, err = IdentityFromEntropy(entropy33)
	if err == nil {
		t.Fatalf("expected error for 33-byte entropy")
	}
}

func TestInvalidXPRV(t *testing.T) {
	// Test with invalid xprv
	_, err := DeriveIdentityFromXPRV("invalid-xprv", 0)
	if err == nil {
		t.Fatalf("expected error for invalid xprv")
	}

	t.Logf("Got expected error for invalid xprv: %v", err)
}