package axon

import (
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha512"
	"errors"
	"fmt"
	"math/big"

	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/curve25519"
)

// p is the prime for Curve25519: 2^255 - 19.
var curveP = new(big.Int).Sub(new(big.Int).Lsh(big.NewInt(1), 255), big.NewInt(19))

// Ed25519PrivKeyToX25519 converts an Ed25519 private key to an X25519 scalar
// by clamping the SHA-512 hash of the 32-byte seed (per RFC 8032 §5.1.5).
func Ed25519PrivKeyToX25519(privKey ed25519.PrivateKey) [32]byte {
	seed := privKey.Seed()
	h := sha512.Sum512(seed)
	// Clamp per RFC 8032.
	h[0] &= 248
	h[31] &= 127
	h[31] |= 64
	var scalar [32]byte
	copy(scalar[:], h[:32])
	return scalar
}

// Ed25519PubKeyToX25519 converts an Ed25519 public key (compressed Edwards
// y-coordinate) to the corresponding Curve25519 Montgomery u-coordinate via
// the birational equivalence: u = (1 + y) / (1 - y) mod p.
func Ed25519PubKeyToX25519(pubKey ed25519.PublicKey) ([32]byte, error) {
	if len(pubKey) != 32 {
		return [32]byte{}, fmt.Errorf("axon: ed25519 pubkey must be 32 bytes, got %d", len(pubKey))
	}

	// Extract the y coordinate: clear the sign bit in byte 31.
	yBytes := make([]byte, 32)
	copy(yBytes, pubKey)
	yBytes[31] &= 0x7f

	// big.Int expects big-endian; reverse the little-endian bytes.
	reversedY := make([]byte, 32)
	for i, b := range yBytes {
		reversedY[31-i] = b
	}
	y := new(big.Int).SetBytes(reversedY)

	p := curveP

	// u = (1 + y) * modInverse(1 - y, p) mod p
	one := big.NewInt(1)

	num := new(big.Int).Add(one, y)
	num.Mod(num, p)

	den := new(big.Int).Sub(one, y)
	den.Mod(den, p)
	den.ModInverse(den, p)

	u := new(big.Int).Mul(num, den)
	u.Mod(u, p)

	// Encode as little-endian 32 bytes.
	uBE := u.Bytes() // big-endian, variable length
	var out [32]byte
	for i, b := range uBE {
		out[len(uBE)-1-i] = b
	}
	return out, nil
}

// DHSharedSecret computes the X25519 Diffie-Hellman shared secret between a
// local Ed25519 private key and a remote Ed25519 public key. Both keys are
// converted to their X25519 equivalents before the scalar multiplication.
func DHSharedSecret(localPriv ed25519.PrivateKey, remotePub ed25519.PublicKey) ([32]byte, error) {
	scalar := Ed25519PrivKeyToX25519(localPriv)
	point, err := Ed25519PubKeyToX25519(remotePub)
	if err != nil {
		return [32]byte{}, fmt.Errorf("axon: convert remote pubkey to x25519: %w", err)
	}
	shared, err := curve25519.X25519(scalar[:], point[:])
	if err != nil {
		return [32]byte{}, fmt.Errorf("axon: x25519: %w", err)
	}
	var out [32]byte
	copy(out[:], shared)
	return out, nil
}

// EncryptDM encrypts plaintext for a DM using ChaCha20-Poly1305. The
// encryption key is the X25519 shared secret derived from senderPriv and
// recipientPub. The returned blob is: nonce (12 bytes) || ciphertext.
func EncryptDM(senderPriv ed25519.PrivateKey, recipientPub ed25519.PublicKey, plaintext []byte) ([]byte, error) {
	shared, err := DHSharedSecret(senderPriv, recipientPub)
	if err != nil {
		return nil, fmt.Errorf("axon: dh: %w", err)
	}

	aead, err := chacha20poly1305.New(shared[:])
	if err != nil {
		return nil, fmt.Errorf("axon: create aead: %w", err)
	}

	nonce := make([]byte, aead.NonceSize()) // 12 bytes
	if _, err := rand.Read(nonce); err != nil {
		return nil, fmt.Errorf("axon: generate nonce: %w", err)
	}

	ct := aead.Seal(nil, nonce, plaintext, nil)
	return append(nonce, ct...), nil
}

// DecryptDM decrypts a blob produced by EncryptDM. recipientPriv is the
// recipient's Ed25519 private key; senderPub is the sender's Ed25519 public
// key. The blob must be nonce (12 bytes) || ciphertext.
func DecryptDM(recipientPriv ed25519.PrivateKey, senderPub ed25519.PublicKey, blob []byte) ([]byte, error) {
	shared, err := DHSharedSecret(recipientPriv, senderPub)
	if err != nil {
		return nil, fmt.Errorf("axon: dh: %w", err)
	}

	aead, err := chacha20poly1305.New(shared[:])
	if err != nil {
		return nil, fmt.Errorf("axon: create aead: %w", err)
	}

	nonceSize := aead.NonceSize()
	if len(blob) < nonceSize {
		return nil, errors.New("axon: ciphertext too short to contain nonce")
	}

	nonce := blob[:nonceSize]
	ct := blob[nonceSize:]

	pt, err := aead.Open(nil, nonce, ct, nil)
	if err != nil {
		return nil, fmt.Errorf("axon: decrypt: authentication failed")
	}
	return pt, nil
}