// Package e2ee is the Go mirror of web/static/crypto.js — the client-side
// encryption every lchat member speaks. It exists so bots can be ordinary Go
// processes instead of Node (docs/bots.md § Direction: a Go sidecar).
//
// The wire format is crypto.js's, byte for byte:
//
//   - AES-256-GCM blobs:   base64url( iv(12) || ciphertext ) — or the raw bytes
//     for attachments.
//   - ECIES to a pubkey:   base64url( ephPub(65) || iv(12) || ciphertext ),
//     key = HKDF-SHA256(ECDH(eph, recipient), salt="", info="lchat-thread-wrap").
//   - Invites:             id / wrapKey / claimSecret each HKDF-derived from the
//     link-fragment secret T with domain-separated info strings.
//
// Compatibility is pinned by testdata/vectors.json, generated from crypto.js
// itself (test/gen-e2ee-vectors.mjs). Any change here or in crypto.js must keep
// those vectors passing — they are the contract.
package e2ee

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"

	"golang.org/x/crypto/hkdf"
)

// B64uEncode matches crypto.js bytesToB64url (unpadded base64url).
func B64uEncode(b []byte) string { return base64.RawURLEncoding.EncodeToString(b) }

// B64uDecode matches crypto.js b64urlToBytes; it tolerates padding.
func B64uDecode(s string) ([]byte, error) {
	for len(s)%4 != 0 {
		s += "="
	}
	return base64.URLEncoding.DecodeString(s)
}

// GenerateKey returns a fresh AES-256 key (raw bytes; B64uEncode to share it).
func GenerateKey() ([]byte, error) {
	key := make([]byte, 32)
	_, err := rand.Read(key)
	return key, err
}

// ImportKey decodes a base64url AES key and checks its length.
func ImportKey(b64 string) ([]byte, error) {
	key, err := B64uDecode(b64)
	if err != nil {
		return nil, err
	}
	if len(key) != 32 {
		return nil, fmt.Errorf("e2ee: AES key is %d bytes, want 32", len(key))
	}
	return key, nil
}

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

// EncryptBytes seals pt and returns raw iv(12)||ciphertext, like crypto.js
// encryptBytes (attachments upload these bytes as octet-stream).
func EncryptBytes(key, pt []byte) ([]byte, error) {
	gcm, err := newGCM(key)
	if err != nil {
		return nil, err
	}
	iv := make([]byte, 12)
	if _, err := rand.Read(iv); err != nil {
		return nil, err
	}
	return gcm.Seal(iv, iv, pt, nil), nil
}

// DecryptBytes opens raw iv(12)||ciphertext.
func DecryptBytes(key, data []byte) ([]byte, error) {
	gcm, err := newGCM(key)
	if err != nil {
		return nil, err
	}
	if len(data) < 12 {
		return nil, fmt.Errorf("e2ee: blob too short")
	}
	return gcm.Open(nil, data[:12], data[12:], nil)
}

// EncryptJSON is crypto.js encrypt: JSON-encode v, seal, base64url the result.
func EncryptJSON(key []byte, v any) (string, error) {
	pt, err := json.Marshal(v)
	if err != nil {
		return "", err
	}
	out, err := EncryptBytes(key, pt)
	if err != nil {
		return "", err
	}
	return B64uEncode(out), nil
}

// DecryptJSON is crypto.js decrypt: open a base64url blob into v.
func DecryptJSON(key []byte, blob string, v any) error {
	data, err := B64uDecode(blob)
	if err != nil {
		return err
	}
	pt, err := DecryptBytes(key, data)
	if err != nil {
		return err
	}
	return json.Unmarshal(pt, v)
}

// ---- invite key-wrapping (see docs/invite-redesign.md) ------------------

// RandomSecret mints a fresh 128-bit invite secret T (base64url).
func RandomSecret() (string, error) {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return B64uEncode(b), nil
}

func hkdfBits(ikm []byte, info string, n int) ([]byte, error) {
	out := make([]byte, n)
	// WebCrypto's empty salt and Go's nil salt both mean zeros[hashLen] (RFC 5869).
	_, err := io.ReadFull(hkdf.New(sha256.New, ikm, nil, []byte(info)), out)
	return out, err
}

// Invite is crypto.js deriveInvite's result: everything derives from T, each
// use domain-separated, so the server (which stores only ID + ClaimHash) can
// neither claim the invite nor unwrap what's wrapped under WrapKey.
type Invite struct {
	ID          string // server lookup id
	WrapKey     []byte // AES key the group key is wrapped under
	ClaimSecret string // presented when claiming; proves knowledge of T
	ClaimHash   string // what the server stores (lowercase hex SHA-256 of ClaimSecret)
}

// DeriveInvite mirrors crypto.js deriveInvite for the secret T (base64url).
func DeriveInvite(tB64 string) (Invite, error) {
	t, err := B64uDecode(tB64)
	if err != nil {
		return Invite{}, err
	}
	id, err := hkdfBits(t, "lchat-invite-id", 32)
	if err != nil {
		return Invite{}, err
	}
	wrapKey, err := hkdfBits(t, "lchat-invite-wrap", 32)
	if err != nil {
		return Invite{}, err
	}
	claim, err := hkdfBits(t, "lchat-invite-claim", 32)
	if err != nil {
		return Invite{}, err
	}
	claimSecret := B64uEncode(claim)
	sum := sha256.Sum256([]byte(claimSecret))
	return Invite{
		ID:          B64uEncode(id),
		WrapKey:     wrapKey,
		ClaimSecret: claimSecret,
		ClaimHash:   hex.EncodeToString(sum[:]),
	}, nil
}

// WrapGroupKey wraps a group key (base64url raw) under an invite's WrapKey.
func WrapGroupKey(groupKeyB64 string, wrapKey []byte) (string, error) {
	raw, err := B64uDecode(groupKeyB64)
	if err != nil {
		return "", err
	}
	out, err := EncryptBytes(wrapKey, raw)
	if err != nil {
		return "", err
	}
	return B64uEncode(out), nil
}

// UnwrapGroupKey recovers the group key (base64url raw).
func UnwrapGroupKey(wrappedB64 string, wrapKey []byte) (string, error) {
	data, err := B64uDecode(wrappedB64)
	if err != nil {
		return "", err
	}
	raw, err := DecryptBytes(wrapKey, data)
	if err != nil {
		return "", err
	}
	return B64uEncode(raw), nil
}

// ---- user keypairs + ECIES (see docs/threads-and-dms.md) ----------------

// Keypair carries a user identity like the browser session does: the public
// key as the raw uncompressed P-256 point (65 bytes), the private key as
// PKCS#8 — both base64url, interchangeable with WebCrypto's raw/pkcs8 exports.
type Keypair struct {
	PublicKey  string
	PrivateKey string
}

// GenerateUserKeypair mirrors crypto.js generateUserKeypair.
func GenerateUserKeypair() (Keypair, error) {
	priv, err := ecdh.P256().GenerateKey(rand.Reader)
	if err != nil {
		return Keypair{}, err
	}
	pkcs8, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return Keypair{}, err
	}
	return Keypair{
		PublicKey:  B64uEncode(priv.PublicKey().Bytes()),
		PrivateKey: B64uEncode(pkcs8),
	}, nil
}

func parsePub(pubB64 string) (*ecdh.PublicKey, error) {
	raw, err := B64uDecode(pubB64)
	if err != nil {
		return nil, err
	}
	return ecdh.P256().NewPublicKey(raw)
}

func parsePriv(privB64 string) (*ecdh.PrivateKey, error) {
	der, err := B64uDecode(privB64)
	if err != nil {
		return nil, err
	}
	key, err := x509.ParsePKCS8PrivateKey(der)
	if err != nil {
		return nil, err
	}
	switch k := key.(type) {
	case *ecdsa.PrivateKey: // NIST-curve PKCS#8 parses as ecdsa; convert
		return k.ECDH()
	case *ecdh.PrivateKey:
		return k, nil
	default:
		return nil, fmt.Errorf("e2ee: unsupported private key type %T", key)
	}
}

// eciesKey mirrors crypto.js eciesKey: ECDH -> HKDF-SHA256 -> AES-256 key.
func eciesKey(priv *ecdh.PrivateKey, pub *ecdh.PublicKey) ([]byte, error) {
	secret, err := priv.ECDH(pub)
	if err != nil {
		return nil, err
	}
	return hkdfBits(secret, "lchat-thread-wrap", 32)
}

// sealTo produces the shared ECIES wire shape: ephPub(65) || iv(12) || ct.
func sealTo(recipientPubB64 string, pt []byte) (string, error) {
	recipient, err := parsePub(recipientPubB64)
	if err != nil {
		return "", err
	}
	eph, err := ecdh.P256().GenerateKey(rand.Reader)
	if err != nil {
		return "", err
	}
	aesKey, err := eciesKey(eph, recipient)
	if err != nil {
		return "", err
	}
	body, err := EncryptBytes(aesKey, pt)
	if err != nil {
		return "", err
	}
	return B64uEncode(append(eph.PublicKey().Bytes(), body...)), nil
}

func openFrom(sealedB64, privKeyB64 string) ([]byte, error) {
	buf, err := B64uDecode(sealedB64)
	if err != nil {
		return nil, err
	}
	if len(buf) < 65+12 {
		return nil, fmt.Errorf("e2ee: sealed blob too short")
	}
	ephPub, err := ecdh.P256().NewPublicKey(buf[:65])
	if err != nil {
		return nil, err
	}
	priv, err := parsePriv(privKeyB64)
	if err != nil {
		return nil, err
	}
	aesKey, err := eciesKey(priv, ephPub)
	if err != nil {
		return nil, err
	}
	return DecryptBytes(aesKey, buf[65:])
}

// WrapThreadKey wraps a thread's AES key (base64url raw) to a recipient's
// public key, exactly like crypto.js wrapThreadKey.
func WrapThreadKey(threadKeyB64, recipientPubB64 string) (string, error) {
	raw, err := B64uDecode(threadKeyB64)
	if err != nil {
		return "", err
	}
	return sealTo(recipientPubB64, raw)
}

// UnwrapThreadKey recovers a thread key (base64url raw) with your private key.
func UnwrapThreadKey(wrappedB64, privKeyB64 string) (string, error) {
	raw, err := openFrom(wrappedB64, privKeyB64)
	if err != nil {
		return "", err
	}
	return B64uEncode(raw), nil
}

// SealForPub seals a JSON-serialisable value to a public key — wrapThreadKey's
// sibling for capability grants (docs/connected-services.md).
func SealForPub(recipientPubB64 string, v any) (string, error) {
	pt, err := json.Marshal(v)
	if err != nil {
		return "", err
	}
	return sealTo(recipientPubB64, pt)
}

// OpenSealed opens a SealForPub blob into v with your private key.
func OpenSealed(sealedB64, privKeyB64 string, v any) error {
	pt, err := openFrom(sealedB64, privKeyB64)
	if err != nil {
		return err
	}
	return json.Unmarshal(pt, v)
}