// Package nip44 implements version 2 of NIP-44, the versioned encrypted
// payload format for Nostr.
//
// A payload is base64(version || nonce || ciphertext || mac), where the
// plaintext has been padded with a length prefix as described in the NIP.
// Conversation keys are derived from a secp256k1 x-only ECDH shared secret
// and all message keys come from HKDF-SHA256.
package nip44

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"math/bits"

	"git.smesh.lol/nostrino/pkg/chacha20"
	"git.smesh.lol/nostrino/pkg/crypto/secp256k1"
)

const (
	// Version is the NIP-44 payload version implemented by this package.
	Version = 2

	nonceSize     = 32
	hmacKeySize   = 32
	chachaKeySize = 32

	// minPayloadChars and minPayloadBytes are the DoS guards required by
	// the NIP: a version byte, a 32-byte nonce, a 32-byte MAC and at least
	// one padded (>= 34 byte) ciphertext block.
	minPayloadChars = 132
	minPayloadBytes = 99

	// minPlaintextSize is 1: empty plaintexts are invalid.
	minPlaintextSize = 1
	// maxPlaintextSize is 2^32 - 1, the largest length the wire format can
	// express.
	maxPlaintextSize = 1<<32 - 1
	// prefixThreshold is the plaintext length at which the padding prefix
	// switches from 2-byte u16 to 6-byte (0x0000 || u32) form.
	prefixThreshold = 1 << 16

	minPaddedSize = 32
)

var (
	errInvalidKey           = errors.New("nip44: invalid secret key or public key")
	errInvalidPlaintextSize = errors.New("nip44: invalid plaintext length")
	errInvalidPayload       = errors.New("nip44: invalid payload")
	errUnknownVersion       = errors.New("nip44: unknown encryption version")
	errInvalidMAC           = errors.New("nip44: invalid MAC")
	errInvalidPadding       = errors.New("nip44: invalid padding")
)

// ConversationKey derives the long-term conversation key shared by two users.
// secret is the local 32-byte secp256k1 private key and peerPub is the peer's
// 32-byte x-only public key. The ECDH result is used unhashed, as required by
// NIP-44 (shared x coordinate -> HKDF-Extract with salt "nip44-v2").
func ConversationKey(secret [32]byte, peerPub [32]byte) ([32]byte, error) {
	var key [32]byte
	shared, ok := secp256k1.ECDH(secret, peerPub)
	if !ok {
		return key, errInvalidKey
	}
	return hkdfExtract([]byte("nip44-v2"), shared[:]), nil
}

// Encrypt encrypts plaintext under conversationKey using a fresh random
// 32-byte nonce and returns the base64 payload.
func Encrypt(conversationKey [32]byte, plaintext []byte) (string, error) {
	var nonce [nonceSize]byte
	if _, err := rand.Read(nonce[:]); err != nil {
		return "", err
	}
	return EncryptWithNonce(conversationKey, nonce, plaintext)
}

// EncryptWithNonce encrypts plaintext under conversationKey using the given
// 32-byte nonce. The nonce must never be reused for the same conversation
// key; callers should normally use Encrypt.
func EncryptWithNonce(conversationKey [32]byte, nonce [32]byte, plaintext []byte) (string, error) {
	padded, err := pad(plaintext)
	if err != nil {
		return "", err
	}
	chachaKey, chachaNonce, hmacKey := messageKeys(conversationKey, nonce)

	ciphertext := make([]byte, len(padded))
	chacha20.XORKeyStream(&chachaKey, &chachaNonce, ciphertext, padded)
	mac := aadMAC(&hmacKey, &nonce, ciphertext)

	raw := make([]byte, 0, 1+nonceSize+len(ciphertext)+32)
	raw = append(raw, Version)
	raw = append(raw, nonce[:]...)
	raw = append(raw, ciphertext...)
	raw = append(raw, mac[:]...)
	return base64.StdEncoding.EncodeToString(raw), nil
}

// Decrypt decrypts a NIP-44 payload using the conversation key. It returns an
// error for any malformed, unauthenticated or wrongly padded payload.
func Decrypt(conversationKey [32]byte, payload string) ([]byte, error) {
	// A leading '#' signals a future, non-base64 encoding.
	if len(payload) == 0 || payload[0] == '#' {
		return nil, errUnknownVersion
	}
	if len(payload) < minPayloadChars {
		return nil, errInvalidPayload
	}
	data, err := base64.StdEncoding.DecodeString(payload)
	if err != nil {
		return nil, errInvalidPayload
	}
	if len(data) < minPayloadBytes {
		return nil, errInvalidPayload
	}
	if data[0] != Version {
		return nil, errUnknownVersion
	}

	var nonce [nonceSize]byte
	copy(nonce[:], data[1:1+nonceSize])
	ciphertextEnd := len(data) - 32
	ciphertext := data[1+nonceSize : ciphertextEnd]
	var mac [32]byte
	copy(mac[:], data[ciphertextEnd:])

	chachaKey, chachaNonce, hmacKey := messageKeys(conversationKey, nonce)
	want := aadMAC(&hmacKey, &nonce, ciphertext)
	if subtle.ConstantTimeCompare(want[:], mac[:]) != 1 {
		return nil, errInvalidMAC
	}

	padded := make([]byte, len(ciphertext))
	chacha20.XORKeyStream(&chachaKey, &chachaNonce, padded, ciphertext)
	return unpad(padded)
}

// hkdfExtract implements HKDF-Extract (RFC 5869) with SHA-256.
func hkdfExtract(salt, ikm []byte) [32]byte {
	mac := hmac.New(sha256.New, salt)
	mac.Write(ikm)
	var prk [32]byte
	mac.Sum(prk[:0])
	return prk
}

// hkdfExpand implements HKDF-Expand (RFC 5869) with SHA-256 for a 32-byte
// PRK. It panics if length exceeds the 255-block limit (8160 bytes); NIP-44
// only ever asks for 76 bytes.
func hkdfExpand(prk [32]byte, info []byte, length int) []byte {
	if length > 255*sha256.Size {
		panic("nip44: hkdf length too large")
	}
	out := make([]byte, 0, length)
	var block []byte // T(0) is empty; each iteration hashes the previous T.
	for counter := 1; len(out) < length; counter++ {
		mac := hmac.New(sha256.New, prk[:])
		mac.Write(block)
		mac.Write(info)
		mac.Write([]byte{byte(counter)})
		block = mac.Sum(block[:0])
		out = append(out, block...)
	}
	return out[:length]
}

// messageKeys derives the ChaCha20 key and nonce and the HMAC key for one
// message from the conversation key and 32-byte nonce.
func messageKeys(conversationKey [32]byte, nonce [32]byte) (key [chachaKeySize]byte, chachaNonce [chacha20.NonceSize]byte, hmacKey [hmacKeySize]byte) {
	okm := hkdfExpand(conversationKey, nonce[:], chachaKeySize+chacha20.NonceSize+hmacKeySize)
	copy(key[:], okm[0:chachaKeySize])
	copy(chachaNonce[:], okm[chachaKeySize:chachaKeySize+chacha20.NonceSize])
	copy(hmacKey[:], okm[chachaKeySize+chacha20.NonceSize:])
	return key, chachaNonce, hmacKey
}

// calcPaddedLen returns the length of the padded plaintext (excluding the
// length prefix) for an unpadded length, as defined by NIP-44.
//
// The calculation deliberately uses 64-bit arithmetic because the result can
// reach 2^32, which overflows a 32-bit integer.
func calcPaddedLen(unpaddedLen uint64) uint64 {
	if unpaddedLen <= minPaddedSize {
		return minPaddedSize
	}
	nextPower := uint64(1) << bits.Len64(unpaddedLen-1)
	var chunk uint64
	if nextPower <= 256 {
		chunk = 32
	} else {
		chunk = nextPower / 8
	}
	return chunk * ((unpaddedLen-1)/chunk + 1)
}

// pad prepends the length prefix and appends zero padding.
func pad(plaintext []byte) ([]byte, error) {
	n := uint64(len(plaintext))
	if n < minPlaintextSize || n > maxPlaintextSize {
		return nil, errInvalidPlaintextSize
	}
	prefixLen := 2
	if n >= prefixThreshold {
		prefixLen = 6
	}
	total := uint64(prefixLen) + calcPaddedLen(n)
	if total > uint64(maxInt()) {
		return nil, errInvalidPlaintextSize
	}
	out := make([]byte, int(total))
	if prefixLen == 6 {
		// The zero u16 signals the extended prefix.
		binary.BigEndian.PutUint32(out[2:6], uint32(n))
	} else {
		binary.BigEndian.PutUint16(out[0:2], uint16(n))
	}
	copy(out[prefixLen:], plaintext)
	return out, nil
}

// unpad validates the length prefix, the padded length and the zero padding,
// then returns the plaintext slice.
func unpad(padded []byte) ([]byte, error) {
	if len(padded) < 2 {
		return nil, errInvalidPadding
	}
	prefixLen := 2
	n := uint64(binary.BigEndian.Uint16(padded[0:2]))
	if n == 0 {
		if len(padded) < 6 {
			return nil, errInvalidPadding
		}
		n = uint64(binary.BigEndian.Uint32(padded[2:6]))
		// A short length in extended form is invalid.
		if n < prefixThreshold {
			return nil, errInvalidPadding
		}
		prefixLen = 6
	}
	if n < minPlaintextSize || n > maxPlaintextSize {
		return nil, errInvalidPadding
	}
	if uint64(len(padded)) != uint64(prefixLen)+calcPaddedLen(n) {
		return nil, errInvalidPadding
	}
	end := prefixLen + int(n)
	for _, b := range padded[end:] {
		if b != 0 {
			return nil, errInvalidPadding
		}
	}
	return padded[prefixLen:end], nil
}

// aadMAC computes HMAC-SHA256 over nonce || ciphertext, as required by
// NIP-44.
func aadMAC(key *[hmacKeySize]byte, nonce *[nonceSize]byte, ciphertext []byte) [32]byte {
	mac := hmac.New(sha256.New, key[:])
	mac.Write(nonce[:])
	mac.Write(ciphertext)
	var out [32]byte
	mac.Sum(out[:0])
	return out
}

// maxInt returns the largest value an int can hold.
func maxInt() int {
	return int(^uint(0) >> 1)
}