package internal

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"io"

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

const (
	saltSize  = 16
	keySize   = 32
	nonceSize = 12
)

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

// Creates a new Curve25519 key pair for identity
func generateIdentityKeyPair() (*IdentityKeyPair, error) {
	var privateKey [32]byte
	if _, err := io.ReadFull(rand.Reader, privateKey[:]); err != nil {
		return nil, err
	}

	var publicKey [32]byte
	curve25519.ScalarBaseMult(&publicKey, &privateKey)

	return &IdentityKeyPair{
		PublicKey:  publicKey,
		PrivateKey: privateKey,
	}, nil
}

// Creates a prekey signed with an Ed25519 signing key
func generateSignedPrekey() ([32]byte, []byte, error) {
	var prekey [32]byte
	var prekeyPriv [32]byte
	if _, err := io.ReadFull(rand.Reader, prekeyPriv[:]); err != nil {
		return prekey, nil, err
	}
	curve25519.ScalarBaseMult(&prekey, &prekeyPriv)

	// Generate signing key
	signingPub, signingPriv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return prekey, nil, err
	}

	// Sign the prekey
	signature := ed25519.Sign(signingPriv, prekey[:])
	combinedSig := append(signingPub, signature...)

	return prekey, combinedSig, nil
}

// Does an ECDH key exchange
func performDH(privateKey, publicKey [32]byte) ([32]byte, error) {
	var sharedSecret [32]byte
	curve25519.ScalarMult(&sharedSecret, &privateKey, &publicKey)
	return sharedSecret, nil
}

// Combines multiple DH outputs to create a shared secret
func deriveSharedSecret(dh1, dh2, dh3 [32]byte) []byte {
	h := sha256.New()
	h.Write(dh1[:])
	h.Write(dh2[:])
	h.Write(dh3[:])
	return h.Sum(nil)
}

// Encrypts a user's private key with their passphrase
func encryptUserKey(privateKey [32]byte, passphrase string) ([]byte, []byte, []byte, error) {
	salt := make([]byte, saltSize)
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		return nil, nil, nil, err
	}

	key := deriveKey(passphrase, salt)
	ciphertext, nonce, err := encryptMsg(string(privateKey[:]), key)
	if err != nil {
		return nil, nil, nil, err
	}

	return ciphertext, nonce, salt, nil
}

// Decrypts a user's private key with their passphrase
func decryptUserKey(ciphertext, nonce, salt []byte, passphrase string) ([32]byte, error) {
	var privateKey [32]byte
	key := deriveKey(passphrase, salt)
	plain, err := decryptMsg(ciphertext, nonce, key)
	if err != nil {
		return privateKey, err
	}
	copy(privateKey[:], []byte(plain))
	return privateKey, nil
}

func deriveKey(passphrase string, salt []byte) []byte {
	return argon2.IDKey([]byte(passphrase), salt, 1, 64*1024, 4, keySize)
}

func encryptMsg(plain string, key []byte) ([]byte, []byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, nil, err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, nil, err
	}
	ciphertext := gcm.Seal(nil, nonce, []byte(plain), nil)
	return ciphertext, nonce, nil
}

func decryptMsg(ciphertext, nonce, key []byte) (string, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}
	plain, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return "", err
	}
	return string(plain), nil
}

// Format: EphemeralPub (32) || Nonce (12) || Ciphertext (...)
func EncryptKeyForUser(recipientPubKey [32]byte, payload []byte) ([]byte, error) {
	var ephPriv, ephPub [32]byte
	if _, err := io.ReadFull(rand.Reader, ephPriv[:]); err != nil {
		return nil, err
	}
	curve25519.ScalarBaseMult(&ephPub, &ephPriv)

	sharedSecret, _ := performDH(ephPriv, recipientPubKey)

	kdf := sha256.Sum256(sharedSecret[:])
	encryptionKey := kdf[:]

	block, err := aes.NewCipher(encryptionKey)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, err
	}

	ciphertext := gcm.Seal(nil, nonce, payload, nil)

	out := make([]byte, 0, 32+len(nonce)+len(ciphertext))
	out = append(out, ephPub[:]...)
	out = append(out, nonce...)
	out = append(out, ciphertext...)

	return out, nil
}

// Decrypts a blob using the user's private identity key.
func DecryptKeyForUser(myPrivKey [32]byte, blob []byte) ([]byte, error) {
	if len(blob) < 32+12 {
		return nil, errors.New("invalid key blob size")
	}

	var ephPub [32]byte
	copy(ephPub[:], blob[:32])
	nonce := blob[32 : 32+12]
	ciphertext := blob[32+12:]

	sharedSecret, _ := performDH(myPrivKey, ephPub)

	kdf := sha256.Sum256(sharedSecret[:])
	decryptionKey := kdf[:]

	block, err := aes.NewCipher(decryptionKey)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	return gcm.Open(nil, nonce, ciphertext, nil)
}