package nostr

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"strings"

	"github.com/btcsuite/btcd/btcec/v2"
)

// NIP04Encrypt encrypts a plaintext message for a recipient using NIP-04
// (ECDH shared secret + AES-256-CBC). The sender's private key and
// recipient's public key (64-char hex) are required.
// Returns the ciphertext in NIP-04 format: "<base64>?iv=<base64>"
func (k *Key) NIP04Encrypt(recipientPubHex, plaintext string) (string, error) {
	if k.priv == nil {
		return "", fmt.Errorf("cannot encrypt: public-only key")
	}

	shared, err := nip04SharedSecret(k.priv, recipientPubHex)
	if err != nil {
		return "", err
	}

	iv := make([]byte, 16)
	if _, err := rand.Read(iv); err != nil {
		return "", fmt.Errorf("generating iv: %w", err)
	}

	block, err := aes.NewCipher(shared)
	if err != nil {
		return "", fmt.Errorf("aes cipher: %w", err)
	}

	padded := pkcs7Pad([]byte(plaintext), aes.BlockSize)
	ct := make([]byte, len(padded))
	cipher.NewCBCEncrypter(block, iv).CryptBlocks(ct, padded)

	return base64.StdEncoding.EncodeToString(ct) + "?iv=" + base64.StdEncoding.EncodeToString(iv), nil
}

// NIP04Decrypt decrypts a NIP-04 ciphertext using the receiver's private key
// and sender's public key (64-char hex).
func (k *Key) NIP04Decrypt(senderPubHex, ciphertext string) (string, error) {
	if k.priv == nil {
		return "", fmt.Errorf("cannot decrypt: public-only key")
	}

	parts := strings.SplitN(ciphertext, "?iv=", 2)
	if len(parts) != 2 {
		return "", fmt.Errorf("invalid NIP-04 ciphertext format")
	}

	ct, err := base64.StdEncoding.DecodeString(parts[0])
	if err != nil {
		return "", fmt.Errorf("decode ciphertext: %w", err)
	}

	iv, err := base64.StdEncoding.DecodeString(parts[1])
	if err != nil {
		return "", fmt.Errorf("decode iv: %w", err)
	}

	shared, err := nip04SharedSecret(k.priv, senderPubHex)
	if err != nil {
		return "", err
	}

	block, err := aes.NewCipher(shared)
	if err != nil {
		return "", fmt.Errorf("aes cipher: %w", err)
	}

	if len(ct)%aes.BlockSize != 0 {
		return "", fmt.Errorf("ciphertext is not a multiple of block size")
	}

	plain := make([]byte, len(ct))
	cipher.NewCBCDecrypter(block, iv).CryptBlocks(plain, ct)

	unpadded, err := pkcs7Unpad(plain)
	if err != nil {
		return "", fmt.Errorf("unpad: %w", err)
	}

	return string(unpadded), nil
}

// nip04SharedSecret computes the ECDH shared secret per NIP-04:
// SHA256(privKey * recipientPubKey).x
func nip04SharedSecret(priv *btcec.PrivateKey, recipientPubHex string) ([]byte, error) {
	pubBytes, err := hex.DecodeString(recipientPubHex)
	if err != nil {
		return nil, fmt.Errorf("decode pubkey: %w", err)
	}

	// NIP-04 pubkeys are x-only (32 bytes, Schnorr). Prefix with 0x02 for
	// compressed SEC1 parsing.
	if len(pubBytes) == 32 {
		pubBytes = append([]byte{0x02}, pubBytes...)
	}

	pub, err := btcec.ParsePubKey(pubBytes)
	if err != nil {
		return nil, fmt.Errorf("parse pubkey: %w", err)
	}

	// ECDH: scalar multiply
	var point btcec.JacobianPoint
	pub.AsJacobian(&point)
	priv.Key.SetByteSlice(priv.Serialize())
	btcec.ScalarMultNonConst(&priv.Key, &point, &point)
	point.ToAffine()

	// Shared secret is SHA256 of the x-coordinate
	xBytes := point.X.Bytes()
	shared := sha256.Sum256(xBytes[:])

	return shared[:], nil
}

func pkcs7Pad(data []byte, blockSize int) []byte {
	pad := blockSize - len(data)%blockSize
	padded := make([]byte, len(data)+pad)
	copy(padded, data)
	for i := len(data); i < len(padded); i++ {
		padded[i] = byte(pad)
	}
	return padded
}

func pkcs7Unpad(data []byte) ([]byte, error) {
	if len(data) == 0 {
		return nil, fmt.Errorf("empty data")
	}
	pad := int(data[len(data)-1])
	if pad == 0 || pad > aes.BlockSize {
		return nil, fmt.Errorf("invalid padding size: %d", pad)
	}
	if len(data) < pad {
		return nil, fmt.Errorf("data shorter than padding")
	}
	return data[:len(data)-pad], nil
}