// Package seal is the mesh's sealed box (novox/hq ADR 0183, ADR 0206): a value sealed to one recipient's
// public key, which only that recipient's private key opens. X25519 for the agreement, HKDF-SHA256 for the
// key, AES-256-GCM for the box — the box the licence manager made and opened before it lived here, byte
// for byte, and the one the SDK's TypeScript `seal` makes and opens.
//
// A box is `{ v: 1, eph, iv, tag, ct }`, every field standard base64; `eph` is the one-time public key as
// SPKI DER. The key is derived with the salt eph-DER ‖ the recipient's raw public key and the info
// "novox-mesh sealed box v1", so a box cannot be re-addressed to another key.
//
// Keys travel as PEM: the public key as SPKI ("PUBLIC KEY"), the private key as PKCS #8 ("PRIVATE KEY").
// Standard library only.
package seal

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
)

// Box is a value sealed to one recipient, as it travels in JSON.
type Box struct {
	V   int    `json:"v"`
	Eph string `json:"eph"`
	IV  string `json:"iv"`
	Tag string `json:"tag"`
	Ct  string `json:"ct"`
}

// KeyPair is a recipient's keypair as the two PEM strings it is kept and sent as.
type KeyPair struct {
	PublicKey  string `json:"publicKey"`
	PrivateKey string `json:"privateKey"`
}

// Version is the only box version Open opens.
const Version = 1

const info = "novox-mesh sealed box v1"

// GenerateKeyPair makes a fresh X25519 keypair, PEM-encoded.
func GenerateKeyPair() (KeyPair, error) {
	priv, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		return KeyPair{}, err
	}
	pubDER, err := x509.MarshalPKIXPublicKey(priv.PublicKey())
	if err != nil {
		return KeyPair{}, err
	}
	privDER, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return KeyPair{}, err
	}
	return KeyPair{
		PublicKey:  string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubDER})),
		PrivateKey: string(pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: privDER})),
	}, nil
}

// Seal seals plaintext to the recipient's PEM public key, with a fresh one-time key and iv.
func Seal(plaintext, recipientPEM string) (Box, error) {
	eph, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		return Box{}, err
	}
	iv := make([]byte, 12)
	if _, err := rand.Read(iv); err != nil {
		return Box{}, err
	}
	return sealWith(plaintext, recipientPEM, eph, iv)
}

// sealWith is Seal with its two sources of randomness given, for the byte-for-byte tests.
func sealWith(plaintext, recipientPEM string, eph *ecdh.PrivateKey, iv []byte) (Box, error) {
	recipient, err := publicFromPEM(recipientPEM)
	if err != nil {
		return Box{}, err
	}
	secret, err := eph.ECDH(recipient)
	if err != nil {
		return Box{}, err
	}
	ephDER, err := x509.MarshalPKIXPublicKey(eph.PublicKey())
	if err != nil {
		return Box{}, err
	}
	gcm, err := newGCM(boxKey(secret, ephDER, recipient.Bytes()))
	if err != nil {
		return Box{}, err
	}
	out := gcm.Seal(nil, iv, []byte(plaintext), nil)
	ct, tag := out[:len(out)-gcm.Overhead()], out[len(out)-gcm.Overhead():]
	b64 := base64.StdEncoding.EncodeToString
	return Box{V: Version, Eph: b64(ephDER), IV: b64(iv), Tag: b64(tag), Ct: b64(ct)}, nil
}

// Open opens a box with the recipient's PEM private key. It fails for a box sealed to another key, one
// changed in any field, or one of another version.
func Open(box Box, privatePEM string) (string, error) {
	if box.V != Version {
		return "", fmt.Errorf("a sealed box of version %d, not %d", box.V, Version)
	}
	priv, err := privateFromPEM(privatePEM)
	if err != nil {
		return "", err
	}
	d := base64.StdEncoding.DecodeString
	ephDER, err := d(box.Eph)
	if err != nil {
		return "", fmt.Errorf("the box's eph: %w", err)
	}
	ephKey, err := x509.ParsePKIXPublicKey(ephDER)
	if err != nil {
		return "", fmt.Errorf("the box's eph: %w", err)
	}
	eph, ok := ephKey.(*ecdh.PublicKey)
	if !ok || eph.Curve() != ecdh.X25519() {
		return "", errors.New("the box's eph is not an X25519 key")
	}
	secret, err := priv.ECDH(eph)
	if err != nil {
		return "", err
	}
	iv, err1 := d(box.IV)
	tag, err2 := d(box.Tag)
	ct, err3 := d(box.Ct)
	if err := errors.Join(err1, err2, err3); err != nil {
		return "", fmt.Errorf("the box's fields: %w", err)
	}
	gcm, err := newGCM(boxKey(secret, ephDER, priv.PublicKey().Bytes()))
	if err != nil {
		return "", err
	}
	if len(iv) != gcm.NonceSize() || len(tag) != gcm.Overhead() {
		return "", errors.New("the box does not open with this key")
	}
	plain, err := gcm.Open(nil, iv, append(append([]byte{}, ct...), tag...), nil)
	if err != nil {
		return "", errors.New("the box does not open with this key")
	}
	return string(plain), nil
}

func publicFromPEM(p string) (*ecdh.PublicKey, error) {
	block, _ := pem.Decode([]byte(p))
	if block == nil {
		return nil, errors.New("not a PEM public key")
	}
	k, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	pub, ok := k.(*ecdh.PublicKey)
	if !ok || pub.Curve() != ecdh.X25519() {
		return nil, errors.New("not an X25519 public key")
	}
	return pub, nil
}

func privateFromPEM(p string) (*ecdh.PrivateKey, error) {
	block, _ := pem.Decode([]byte(p))
	if block == nil {
		return nil, errors.New("not a PEM private key")
	}
	k, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	priv, ok := k.(*ecdh.PrivateKey)
	if !ok || priv.Curve() != ecdh.X25519() {
		return nil, errors.New("not an X25519 private key")
	}
	return priv, nil
}

// boxKey is HKDF-SHA256 (RFC 5869) for one 32-byte block: written out because crypto/hkdf needs Go 1.24
// and this module promises 1.22; the vectors prove it equals the licence manager's hkdf.Key.
func boxKey(secret, ephDER, recipientRaw []byte) []byte {
	salt := append(append([]byte{}, ephDER...), recipientRaw...)
	extract := hmac.New(sha256.New, salt)
	extract.Write(secret)
	prk := extract.Sum(nil)
	expand := hmac.New(sha256.New, prk)
	expand.Write([]byte(info))
	expand.Write([]byte{1})
	return expand.Sum(nil)
}

func newGCM(key []byte) (cipher.AEAD, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return cipher.NewGCM(block)
}