package x25519

import (
	"bytes"
	"crypto/ecdh"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"io"
	"testing"
)

func TestGenerateKeyPair(t *testing.T) {
	adapter := New()

	pub, priv, err := adapter.GenerateKeyPair(rand.Reader)
	if err != nil {
		t.Fatalf("GenerateKeyPair failed: %v", err)
	}

	if pub.SchemeID() != suiteID {
		t.Errorf("pub.SchemeID() = 0x%04x, want 0x%04x", pub.SchemeID(), suiteID)
	}

	if priv.SchemeID() != suiteID {
		t.Errorf("priv.SchemeID() = 0x%04x, want 0x%04x", priv.SchemeID(), suiteID)
	}

	rawPub := pub.Raw()
	if len(rawPub) != 32 {
		t.Errorf("pub.Raw() len = %d, want 32", len(rawPub))
	}

	rawPriv := priv.Raw()
	if len(rawPriv) != 32 {
		t.Errorf("priv.Raw() len = %d, want 32", len(rawPriv))
	}

	expectedKeyID := sha256.Sum256(rawPub)
	if !bytes.Equal(pub.KeyID(), expectedKeyID[:8]) {
		t.Errorf("pub.KeyID() = %x, want %x", pub.KeyID(), expectedKeyID[:8])
	}

	if !bytes.Equal(priv.KeyID(), pub.KeyID()) {
		t.Errorf("priv.KeyID() = %x, want %x", priv.KeyID(), pub.KeyID())
	}
}

func TestEncapsulate(t *testing.T) {
	adapter := New()

	pub, priv, err := adapter.GenerateKeyPair(rand.Reader)
	if err != nil {
		t.Fatalf("GenerateKeyPair failed: %v", err)
	}

	ct, ss, err := adapter.Encapsulate(pub, rand.Reader)
	if err != nil {
		t.Fatalf("Encapsulate failed: %v", err)
	}

	if len(ct) != 32 {
		t.Errorf("ciphertext len = %d, want 32", len(ct))
	}

	if len(ss) != 32 {
		t.Errorf("sharedSecret len = %d, want 32", len(ss))
	}

	// Verify ss by independently computing priv.ECDH(ephemeralPubParsedFromCt).
	ephPub, err := ecdh.X25519().NewPublicKey(ct)
	if err != nil {
		t.Fatalf("failed to parse ephemeral public key from ciphertext: %v", err)
	}

	parsedPriv, err := ecdh.X25519().NewPrivateKey(priv.Raw())
	if err != nil {
		t.Fatalf("failed to parse private key: %v", err)
	}

	computedSS, err := parsedPriv.ECDH(ephPub)
	if err != nil {
		t.Fatalf("independent ECDH computation failed: %v", err)
	}

	if !bytes.Equal(ss, computedSS) {
		t.Errorf("shared secret mismatch: encapsulate=%x, independent=%x", ss, computedSS)
	}
}

func TestRoundTrip(t *testing.T) {
	adapter := New()

	pub, priv, err := adapter.GenerateKeyPair(rand.Reader)
	if err != nil {
		t.Fatalf("GenerateKeyPair failed: %v", err)
	}

	ct, ssEnc, err := adapter.Encapsulate(pub, rand.Reader)
	if err != nil {
		t.Fatalf("Encapsulate failed: %v", err)
	}

	ssDec, err := adapter.Decapsulate(priv, ct)
	if err != nil {
		t.Fatalf("Decapsulate failed: %v", err)
	}

	if !bytes.Equal(ssEnc, ssDec) {
		t.Fatalf("shared secret mismatch: encapsulate=%x, decapsulate=%x", ssEnc, ssDec)
	}
}

func TestRoundTripMany(t *testing.T) {
	adapter := New()

	for iteration := 0; iteration < 1000; iteration++ {
		pub, priv, err := adapter.GenerateKeyPair(rand.Reader)
		if err != nil {
			t.Fatalf("iteration %d: GenerateKeyPair failed: %v", iteration, err)
		}

		ct, ssEnc, err := adapter.Encapsulate(pub, rand.Reader)
		if err != nil {
			t.Fatalf("iteration %d: Encapsulate failed: %v", iteration, err)
		}

		ssDec, err := adapter.Decapsulate(priv, ct)
		if err != nil {
			t.Fatalf("iteration %d: Decapsulate failed: %v", iteration, err)
		}

		if !bytes.Equal(ssEnc, ssDec) {
			t.Fatalf("iteration %d: shared secret mismatch", iteration)
		}
	}
}

func TestDecapsulateRandomCiphertext(t *testing.T) {
	adapter := New()

	_, priv, err := adapter.GenerateKeyPair(rand.Reader)
	if err != nil {
		t.Fatalf("GenerateKeyPair failed: %v", err)
	}

	// Generate a random 32-byte string that is unlikely to be a valid X25519 public key.
	// Setting the high bit makes it invalid for X25519 (Montgomery u-coordinate must be < 2^255).
	randomCT := make([]byte, 32)
	if _, err := io.ReadFull(rand.Reader, randomCT); err != nil {
		t.Fatalf("failed to read random bytes: %v", err)
	}
	randomCT[31] |= 0x80 // set high bit to guarantee invalidity

	_, err = adapter.Decapsulate(priv, randomCT)
	if err == nil {
		t.Fatal("Decapsulate with random ciphertext: expected error, got nil")
	}

	if !errors.Is(err, ErrDecapsulationFailed) {
		t.Errorf("Decapsulate error = %v, want ErrDecapsulationFailed", err)
	}
}

func TestRegistryRegistration(t *testing.T) {
	factory, err := DefaultRegistry.Lookup(suiteID)
	if err != nil {
		t.Fatalf("Lookup suiteID 0x%04x failed: %v", suiteID, err)
	}

	instance := factory()
	if instance.SchemeID() != suiteID {
		t.Errorf("factory() SchemeID = 0x%04x, want 0x%04x", instance.SchemeID(), suiteID)
	}
}

func TestFactoryReturnsIndependentInstances(t *testing.T) {
	factory, err := DefaultRegistry.Lookup(suiteID)
	if err != nil {
		t.Fatalf("Lookup suiteID 0x%04x failed: %v", suiteID, err)
	}

	one := factory()
	two := factory()

	if one.SchemeID() != two.SchemeID() {
		t.Error("factory() returned instances with different scheme IDs")
	}
}
