package x25519

import (
	"crypto/ecdh"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"

	"git.tswf.io/infra/go-synapse-backupper/pkg/domain/crypto"
)

// suiteID is the scheme identifier for X25519 ECDH KEM.
const suiteID uint16 = 0x0007

// ErrDecapsulationFailed is returned when ciphertext decapsulation fails,
// typically because the ciphertext is not a valid X25519 public key.
var ErrDecapsulationFailed = errors.New("decapsulation failed")

// DefaultRegistry is the package-level registry for X25519.
var DefaultRegistry = crypto.NewRegistry()

// kemAdapter wraps the Go stdlib crypto/ecdh X25519 implementation to satisfy
// the pkg/domain/crypto.KEM interface.
type kemAdapter struct{}

// New creates a new KEM adapter instance.
func New() crypto.KEM {
	return &kemAdapter{}
}

// SchemeID returns the X25519 scheme identifier (0x0007).
func (k *kemAdapter) SchemeID() uint16 {
	return suiteID
}

// GenerateKeyPair generates a new X25519 key pair.
func (k *kemAdapter) GenerateKeyPair(
	rand io.Reader,
) (
	crypto.RecipientPub,
	crypto.RecipientPriv,
	error,
) {
	ecdhPriv, err := ecdh.X25519().GenerateKey(rand)
	if err != nil {
		return nil, nil, err
	}

	rawPub := ecdhPriv.PublicKey().Bytes()
	keyID := computeKeyID(rawPub)

	pub := &pubKey{
		key:   ecdhPriv.PublicKey(),
		keyID: keyID,
	}
	priv := &privKey{
		key:   ecdhPriv,
		keyID: keyID,
	}

	return pub, priv, nil
}

// Encapsulate generates a shared secret and ciphertext for the given public key.
// The ciphertext is the ephemeral public key (32 bytes).
func (k *kemAdapter) Encapsulate(
	pub crypto.RecipientPub,
	rand io.Reader,
) (
	ciphertext []byte,
	sharedSecret []byte,
	err error,
) {
	p, ok := pub.(*pubKey)
	if !ok {
		raw := pub.Raw()
		ek, parseErr := ecdh.X25519().NewPublicKey(raw)
		if parseErr != nil {
			return nil, nil, fmt.Errorf("invalid public key for X25519: %w", parseErr)
		}
		p = &pubKey{key: ek, keyID: computeKeyID(raw)}
	}

	ephPriv, err := ecdh.X25519().GenerateKey(rand)
	if err != nil {
		return nil, nil, err
	}

	ct := ephPriv.PublicKey().Bytes()
	ss, err := ephPriv.ECDH(p.key)
	if err != nil {
		return nil, nil, err
	}

	return ct, ss, nil
}

// LoadPriv loads an X25519 private key from raw bytes.
func (k *kemAdapter) LoadPriv(
	raw []byte,
) (
	crypto.RecipientPriv,
	error,
) {
	dk, err := ecdh.X25519().NewPrivateKey(raw)
	if err != nil {
		return nil, fmt.Errorf("invalid X25519 private key: %w", err)
	}
	pubRaw := dk.PublicKey().Bytes()
	return &privKey{key: dk, keyID: computeKeyID(pubRaw)}, nil
}

// Decapsulate recovers the shared secret from a ciphertext using the private key.
// The ciphertext must be a valid 32-byte X25519 public key.
func (k *kemAdapter) Decapsulate(
	priv crypto.RecipientPriv,
	ciphertext []byte,
) (
	sharedSecret []byte,
	err error,
) {
	p, ok := priv.(*privKey)
	if !ok {
		return nil, fmt.Errorf("invalid private key type for X25519")
	}

	if len(ciphertext) != 32 {
		return nil, errors.Join(ErrDecapsulationFailed, errors.New("invalid ciphertext length"))
	}

	// X25519 public keys are 255-bit Montgomery u-coordinates; bit 255 must be zero.
	if ciphertext[31]&0x80 != 0 {
		return nil, ErrDecapsulationFailed
	}

	// Reject the all-zero public key (identity point), which yields an all-zero shared secret.
	allZero := true
	for _, b := range ciphertext {
		if b != 0 {
			allZero = false
			break
		}
	}
	if allZero {
		return nil, ErrDecapsulationFailed
	}

	ephPub, err := ecdh.X25519().NewPublicKey(ciphertext)
	if err != nil {
		return nil, errors.Join(ErrDecapsulationFailed, err)
	}

	ss, err := p.key.ECDH(ephPub)
	if err != nil {
		return nil, errors.Join(ErrDecapsulationFailed, err)
	}

	return ss, nil
}

// pubKey wraps *ecdh.PublicKey to satisfy crypto.RecipientPub.
type pubKey struct {
	key   *ecdh.PublicKey
	keyID []byte
}

func (p *pubKey) SchemeID() uint16 { return suiteID }
func (p *pubKey) KeyID() []byte    { return p.keyID }
func (p *pubKey) Raw() []byte      { return p.key.Bytes() }

// privKey wraps *ecdh.PrivateKey to satisfy crypto.RecipientPriv.
type privKey struct {
	key   *ecdh.PrivateKey
	keyID []byte
}

func (p *privKey) SchemeID() uint16 { return suiteID }
func (p *privKey) KeyID() []byte    { return p.keyID }
func (p *privKey) Raw() []byte      { return p.key.Bytes() }

// computeKeyID derives the first 8 bytes of SHA-256 over the raw public key.
func computeKeyID(raw []byte) []byte {
	h := sha256.Sum256(raw)
	return h[:8]
}

// init registers the X25519 factory under suiteID 0x0007.
func init() {
	_ = DefaultRegistry.Register(
		suiteID,
		func() crypto.KEM {
			return New()
		},
	)
}