package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/hkdf"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"log"
	"math/big"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

// Web Push, implemented against the RFC rather than pulled in as a dependency:
// message encryption is RFC 8291 (ECDH to a per-message key) wrapped in the
// RFC 8188 aes128gcm content encoding, and the request is authorized with a
// VAPID (RFC 8292) ES256 JWT identifying this server to the push service.
// It is ~150 lines of stdlib crypto, and encryptPayload is checked against the
// RFC 8291 §5 test vector in webpush_test.go.

// b64 is the unpadded base64url alphabet every web push field uses: the keys a
// browser hands us in a PushSubscription, the JWT segments, and the VAPID key.
var b64 = base64.RawURLEncoding

// Subscription is a browser's PushSubscription: where to send, plus the two
// keys its service worker will decrypt with. Stored verbatim per device.
type Subscription struct {
	Endpoint string `json:"endpoint"`
	Keys     struct {
		P256dh string `json:"p256dh"` // the client's public key, uncompressed P-256 point
		Auth   string `json:"auth"`   // 16-byte shared authentication secret
	} `json:"keys"`
}

// VAPIDKey is this server's identity to push services. The same key must be
// used for the lifetime of a subscription: browsers pin the public key given at
// subscribe time, so rotating it invalidates every existing subscription.
type VAPIDKey struct {
	priv *ecdsa.PrivateKey
	// Public is the uncompressed public point, base64url — handed to the client
	// as applicationServerKey and echoed in the Authorization header.
	Public string
}

// vapidFile is the on-disk form of a VAPID key: just the P-256 scalar, so the
// public half is always rederived and can never drift out of sync with it.
type vapidFile struct {
	Private string `json:"private"`
}

func newVAPIDKey() (*VAPIDKey, error) {
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}
	return vapidFromKey(priv), nil
}

func vapidFromKey(priv *ecdsa.PrivateKey) *VAPIDKey {
	pub, _ := priv.PublicKey.ECDH()
	return &VAPIDKey{priv: priv, Public: b64.EncodeToString(pub.Bytes())}
}

func (k *VAPIDKey) marshal() ([]byte, error) {
	return json.MarshalIndent(vapidFile{Private: b64.EncodeToString(k.priv.D.FillBytes(make([]byte, 32)))}, "", "  ")
}

func parseVAPIDKey(raw []byte) (*VAPIDKey, error) {
	var f vapidFile
	if err := json.Unmarshal(raw, &f); err != nil {
		return nil, err
	}
	return vapidFromSeed(f.Private)
}

// vapidFromSeed rebuilds the keypair from the base64url private scalar, the form
// both the key file and the -vapid-key flag carry.
func vapidFromSeed(seed string) (*VAPIDKey, error) {
	d, err := b64.DecodeString(strings.TrimSpace(seed))
	if err != nil {
		return nil, fmt.Errorf("decode vapid key: %w", err)
	}
	if len(d) != 32 {
		return nil, fmt.Errorf("vapid key must be 32 bytes, got %d", len(d))
	}
	ecdhPriv, err := ecdh.P256().NewPrivateKey(d)
	if err != nil {
		return nil, fmt.Errorf("invalid vapid key: %w", err)
	}
	// crypto/ecdh validated the scalar and derived the point for us; split the
	// uncompressed encoding (0x04 | X | Y) back into the coordinates ecdsa wants.
	point := ecdhPriv.PublicKey().Bytes()
	if len(point) != 65 || point[0] != 4 {
		return nil, fmt.Errorf("invalid vapid key: bad public point")
	}
	priv := &ecdsa.PrivateKey{
		PublicKey: ecdsa.PublicKey{
			Curve: elliptic.P256(),
			X:     new(big.Int).SetBytes(point[1:33]),
			Y:     new(big.Int).SetBytes(point[33:]),
		},
		D: new(big.Int).SetBytes(d),
	}
	return vapidFromKey(priv), nil
}

// authHeader builds the VAPID Authorization header for one push endpoint. The
// audience is the endpoint's origin — a token minted for one push service is
// not valid at another — and the short expiry bounds replay if it leaks.
func (k *VAPIDKey) authHeader(endpoint, subject string) (string, error) {
	u, err := url.Parse(endpoint)
	if err != nil {
		return "", err
	}
	claims := map[string]any{
		"aud": u.Scheme + "://" + u.Host,
		"exp": time.Now().Add(12 * time.Hour).Unix(),
		"sub": subject,
	}
	body, err := json.Marshal(claims)
	if err != nil {
		return "", err
	}
	// Header is constant for ES256, so it is spelled out rather than marshalled.
	signing := b64.EncodeToString([]byte(`{"typ":"JWT","alg":"ES256"}`)) + "." + b64.EncodeToString(body)
	sum := sha256.Sum256([]byte(signing))
	r, s, err := ecdsa.Sign(rand.Reader, k.priv, sum[:])
	if err != nil {
		return "", err
	}
	// JWS wants the raw r||s pair, fixed-width — not the ASN.1 sequence
	// ecdsa.SignASN1 would give us.
	sig := make([]byte, 64)
	r.FillBytes(sig[:32])
	s.FillBytes(sig[32:])
	jwt := signing + "." + b64.EncodeToString(sig)
	return "vapid t=" + jwt + ", k=" + k.Public, nil
}

// encryptPayload encrypts plaintext for one subscription per RFC 8291, emitting
// a complete RFC 8188 aes128gcm body: a header carrying the salt and this
// message's ephemeral public key, followed by a single AES-GCM record.
//
// salt and the ephemeral key are parameters rather than generated inline purely
// so the RFC test vector can be reproduced; callers pass nil for both.
func encryptPayload(sub Subscription, plaintext, salt []byte, eph *ecdh.PrivateKey) ([]byte, error) {
	clientPubRaw, err := b64.DecodeString(sub.Keys.P256dh)
	if err != nil {
		return nil, fmt.Errorf("decode p256dh: %w", err)
	}
	authSecret, err := b64.DecodeString(sub.Keys.Auth)
	if err != nil {
		return nil, fmt.Errorf("decode auth: %w", err)
	}
	clientPub, err := ecdh.P256().NewPublicKey(clientPubRaw)
	if err != nil {
		return nil, fmt.Errorf("invalid p256dh: %w", err)
	}
	if eph == nil {
		if eph, err = ecdh.P256().GenerateKey(rand.Reader); err != nil {
			return nil, err
		}
	}
	if salt == nil {
		salt = make([]byte, 16)
		if _, err := rand.Read(salt); err != nil {
			return nil, err
		}
	}
	shared, err := eph.ECDH(clientPub)
	if err != nil {
		return nil, fmt.Errorf("ecdh: %w", err)
	}

	// RFC 8291 §3.4: the auth secret salts a first extraction that binds the
	// derived key to *both* public keys, so a message can only be decrypted by
	// the subscription it was addressed to.
	ephPub := eph.PublicKey().Bytes()
	keyInfo := append([]byte("WebPush: info\x00"), clientPubRaw...)
	keyInfo = append(keyInfo, ephPub...)
	ikm, err := hkdf.Key(sha256.New, shared, authSecret, string(keyInfo), 32)
	if err != nil {
		return nil, err
	}
	cek, err := hkdf.Key(sha256.New, ikm, salt, "Content-Encoding: aes128gcm\x00", 16)
	if err != nil {
		return nil, err
	}
	nonce, err := hkdf.Key(sha256.New, ikm, salt, "Content-Encoding: nonce\x00", 12)
	if err != nil {
		return nil, err
	}

	block, err := aes.NewCipher(cek)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	// Single record, so the padding delimiter is 0x02 ("last record") with no
	// padding after it. Multi-record chunking would use 0x01 for earlier records.
	record := gcm.Seal(nil, nonce, append(append([]byte{}, plaintext...), 0x02), nil)

	// RFC 8188 §2.1 header: salt | record size | key id length | key id.
	var out bytes.Buffer
	out.Write(salt)
	_ = binary.Write(&out, binary.BigEndian, uint32(4096))
	out.WriteByte(byte(len(ephPub)))
	out.Write(ephPub)
	out.Write(record)
	return out.Bytes(), nil
}

// pushError reports a push service rejecting a send. Gone is set for the 404 and
// 410 responses that mean the subscription is permanently dead, which is the
// signal callers use to drop it — any other failure is transient and kept.
type pushError struct {
	Status int
	Body   string
	Gone   bool
}

func (e *pushError) Error() string {
	return fmt.Sprintf("push service returned %d: %s", e.Status, e.Body)
}

// send delivers one encrypted message to a subscription's endpoint.
func (k *VAPIDKey) send(client *http.Client, sub Subscription, payload []byte, ttl int) error {
	body, err := encryptPayload(sub, payload, nil, nil)
	if err != nil {
		return err
	}
	auth, err := k.authHeader(sub.Endpoint, vapidSubject)
	if err != nil {
		return err
	}
	req, err := http.NewRequest(http.MethodPost, sub.Endpoint, bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", auth)
	req.Header.Set("Content-Encoding", "aes128gcm")
	req.Header.Set("Content-Type", "application/octet-stream")
	req.Header.Set("TTL", fmt.Sprint(ttl))
	// Reminders are only useful while current: if the device is offline long
	// enough for a later evaluation to supersede this one, dropping it is right.
	req.Header.Set("Urgency", "normal")

	res, err := client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode >= 200 && res.StatusCode < 300 {
		return nil
	}
	msg := make([]byte, 512)
	n, _ := res.Body.Read(msg)
	return &pushError{
		Status: res.StatusCode,
		Body:   strings.TrimSpace(string(msg[:n])),
		Gone:   res.StatusCode == http.StatusNotFound || res.StatusCode == http.StatusGone,
	}
}

// vapidSubject identifies this server to push services. RFC 8292 wants a
// contact URL; push services in practice only require that it be present and
// well-formed, and this app has no operator address to offer.
const vapidSubject = "mailto:puppy-tracker@localhost"

// loadVAPIDKey resolves the server's push identity. An explicit seed (flag or
// env) wins so deployments can hold the key in a secrets file; otherwise it is
// read from path, and generated and persisted there on first run. Rotating this
// key silently breaks every existing subscription, so it is only ever created
// when absent — never regenerated on a read error.
func loadVAPIDKey(seed, path string) (*VAPIDKey, error) {
	if seed != "" {
		return vapidFromSeed(seed)
	}
	raw, err := os.ReadFile(path)
	if err == nil {
		return parseVAPIDKey(raw)
	}
	if !os.IsNotExist(err) {
		return nil, err
	}
	key, err := newVAPIDKey()
	if err != nil {
		return nil, err
	}
	out, err := key.marshal()
	if err != nil {
		return nil, err
	}
	// 0600: the private half is the only thing stopping someone else pushing
	// notifications to this app's users.
	if err := os.WriteFile(path, out, 0o600); err != nil {
		return nil, err
	}
	log.Printf("generated VAPID key at %s", path)
	return key, nil
}