package crypto

import (
	"crypto/rand"
	"crypto/sha256"
	"io"

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

type E2EEKeypair struct {
	PrivateKey [32]byte
	PublicKey  [32]byte
}

func GenerateE2EEKeypair() (*E2EEKeypair, error) {
	priv := make([]byte, 32)
	if _, err := io.ReadFull(rand.Reader, priv); err != nil {
		return nil, err
	}
	// Clamp for X25519
	priv[0] &= 248
	priv[31] &= 127
	priv[31] |= 64

	pub, err := curve25519.X25519(priv, curve25519.Basepoint)
	if err != nil {
		return nil, err
	}
	var privArr [32]byte
	var pubArr [32]byte
	copy(privArr[:], priv)
	copy(pubArr[:], pub)
	return &E2EEKeypair{PrivateKey: privArr, PublicKey: pubArr}, nil
}

func DeriveE2EEKey(privateKey, publicKey [32]byte) []byte {
	shared, _ := curve25519.X25519(privateKey[:], publicKey[:])

	salt := make([]byte, 32)
	info := []byte("qcc-e2ee-media-key")
	hkdf := hkdf.New(sha256.New, shared, salt, info)
	key := make([]byte, 32)
	if _, err := io.ReadFull(hkdf, key); err != nil {
		return nil
	}
	return key
}

func EncryptE2EE(key []byte, plaintext []byte) ([]byte, error) {
	aead, err := chacha20poly1305.New(key)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, chacha20poly1305.NonceSizeX)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, err
	}
	return aead.Seal(nonce, nonce, plaintext, nil), nil
}

func DecryptE2EE(key []byte, ciphertext []byte) ([]byte, error) {
	aead, err := chacha20poly1305.New(key)
	if err != nil {
		return nil, err
	}
	if len(ciphertext) < chacha20poly1305.NonceSizeX {
		return nil, io.ErrUnexpectedEOF
	}
	nonce := ciphertext[:chacha20poly1305.NonceSizeX]
	return aead.Open(nil, nonce, ciphertext[chacha20poly1305.NonceSizeX:], nil)
}