package internal

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

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

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

// IdentityKeyPair represents a user's long-term identity key
type IdentityKeyPair struct {
	PublicKey  [32]byte
	PrivateKey [32]byte
}

// PrekeyBundle contains public keys for establishing a session
type PrekeyBundle struct {
	IdentityKey     [32]byte
	Prekey          [32]byte
	PrekeySignature []byte
}

// generateIdentityKeyPair 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
}

// generateSignedPrekey 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[:])

	// In practice, we'd store signingPub to verify later. For now, we include it in signature
	combinedSig := append(signingPub, signature...)

	return prekey, combinedSig, nil
}

// performDH 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
}

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

// encryptUserKey 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
}

// decryptUserKey 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[:], 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
}