package crypto

import (
	"errors"
	"io"
	"sync"
	"testing"
)

// mockKEM is a minimal KEM implementation for testing the registry.
type mockKEM struct {
	schemeID uint16
}

func (m *mockKEM) SchemeID() uint16 {
	return m.schemeID
}

func (m *mockKEM) GenerateKeyPair(
	rand io.Reader,
) (
	RecipientPub,
	RecipientPriv,
	error,
) {
	return nil, nil, nil
}

func (m *mockKEM) Encapsulate(
	pub RecipientPub,
	rand io.Reader,
) (
	ciphertext []byte,
	sharedSecret []byte,
	err error,
) {
	return nil, nil, nil
}

func (m *mockKEM) Decapsulate(
	priv RecipientPriv,
	ciphertext []byte,
) (
	sharedSecret []byte,
	err error,
) {
	return nil, nil
}

func (m *mockKEM) LoadPriv(
	raw []byte,
) (
	RecipientPriv,
	error,
) {
	return &mockRecipient{schemeIDValue: m.schemeID, rawBytes: raw}, nil
}

// mockRecipient implements RecipientPub and RecipientPriv for tests.
type mockRecipient struct {
	schemeIDValue uint16
	rawBytes      []byte
}

func (r *mockRecipient) SchemeID() uint16 {
	return r.schemeIDValue
}

func (r *mockRecipient) KeyID() []byte {
	return r.rawBytes
}

func (r *mockRecipient) Raw() []byte {
	return r.rawBytes
}

func TestRegistry_RegisterAndLookup(t *testing.T) {
	registry := NewRegistry()
	factory := func() KEM {
		return &mockKEM{schemeID: 0x0001}
	}

	err := registry.Register(0x0001, factory)
	if err != nil {
		t.Fatalf("unexpected error registering scheme: %v", err)
	}

	foundFactory, err := registry.Lookup(0x0001)
	if err != nil {
		t.Fatalf("unexpected error looking up scheme: %v", err)
	}

	kem := foundFactory()
	if kem.SchemeID() != 0x0001 {
		t.Errorf("expected schemeID 0x0001, got 0x%04x", kem.SchemeID())
	}
}

func TestRegistry_LookupUnknownScheme(t *testing.T) {
	registry := NewRegistry()

	_, err := registry.Lookup(0x9999)
	if err == nil {
		t.Fatal("expected error for unknown scheme, got nil")
	}

	if !errors.Is(err, ErrUnknownScheme) {
		t.Errorf("expected ErrUnknownScheme, got %v", err)
	}
}

func TestRegistry_RegisterDuplicateScheme(t *testing.T) {
	registry := NewRegistry()
	factory := func() KEM {
		return &mockKEM{schemeID: 0x0001}
	}

	err := registry.Register(0x0001, factory)
	if err != nil {
		t.Fatalf("unexpected error on first register: %v", err)
	}

	err = registry.Register(0x0001, factory)
	if err == nil {
		t.Fatal("expected error for duplicate scheme, got nil")
	}

	if !errors.Is(err, ErrDuplicateScheme) {
		t.Errorf("expected ErrDuplicateScheme, got %v", err)
	}
}

func TestRegistry_ConcurrentAccess(t *testing.T) {
	registry := NewRegistry()
	const goroutines = 100

	var waitGroup sync.WaitGroup
	waitGroup.Add(goroutines)

	for index := range goroutines {
		go func(schemeID uint16) {
			defer waitGroup.Done()

			factory := func() KEM {
				return &mockKEM{schemeID: schemeID}
			}

			_ = registry.Register(schemeID, factory)
			_, _ = registry.Lookup(schemeID)
		}(uint16(index + 1))
	}

	waitGroup.Wait()
}

func TestRegistry_LookupReturnsIndependentInstances(t *testing.T) {
	registry := NewRegistry()
	factory := func() KEM {
		return &mockKEM{schemeID: 0x0001}
	}

	_ = registry.Register(0x0001, factory)

	factoryOne, _ := registry.Lookup(0x0001)
	factoryTwo, _ := registry.Lookup(0x0001)

	kemOne := factoryOne()
	kemTwo := factoryTwo()

	if kemOne == kemTwo {
		t.Error("expected independent KEM instances from factory")
	}
}