package sandbar

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/hkdf"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"io"
	"math/big"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"sync"
	"testing"
	"time"
)

// fakeDevice plays the browser: it owns the subscription keys and can decrypt what arrives.
type fakeDevice struct {
	priv   *ecdh.PrivateKey
	auth   []byte
	srv    *httptest.Server
	mu     sync.Mutex
	got    []pushPayload
	vapid  string
	status int
}

func newFakeDevice(t *testing.T) *fakeDevice {
	t.Helper()
	priv, _ := ecdh.P256().GenerateKey(rand.Reader)
	d := &fakeDevice{priv: priv, auth: make([]byte, 16), status: 201}
	rand.Read(d.auth)
	d.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		d.mu.Lock()
		defer d.mu.Unlock()
		d.vapid = r.Header.Get("Authorization")
		if d.status != 201 {
			w.WriteHeader(d.status)
			return
		}
		pt, err := d.decrypt(body)
		if err != nil {
			t.Errorf("decrypt: %v", err)
			w.WriteHeader(400)
			return
		}
		var p pushPayload
		json.Unmarshal(pt, &p)
		d.got = append(d.got, p)
		w.WriteHeader(201)
	}))
	t.Cleanup(d.srv.Close)
	return d
}

func (d *fakeDevice) keys() map[string]string {
	return map[string]string{
		"p256dh": base64.RawURLEncoding.EncodeToString(d.priv.PublicKey().Bytes()),
		"auth":   base64.RawURLEncoding.EncodeToString(d.auth),
	}
}

// decrypt is RFC 8291 from the receiving side, written independently of encryptPush.
func (d *fakeDevice) decrypt(b []byte) ([]byte, error) {
	salt, idlen := b[:16], int(b[20])
	asPubRaw := b[21 : 21+idlen]
	ct := b[21+idlen:]
	asPub, err := ecdh.P256().NewPublicKey(asPubRaw)
	if err != nil {
		return nil, err
	}
	secret, _ := d.priv.ECDH(asPub)
	prkKey, _ := hkdf.Extract(sha256.New, secret, d.auth)
	info := append(append([]byte("WebPush: info\x00"), d.priv.PublicKey().Bytes()...), asPubRaw...)
	ikm, _ := hkdf.Expand(sha256.New, prkKey, string(info), 32)
	prk, _ := hkdf.Extract(sha256.New, ikm, salt)
	cek, _ := hkdf.Expand(sha256.New, prk, "Content-Encoding: aes128gcm\x00", 16)
	nonce, _ := hkdf.Expand(sha256.New, prk, "Content-Encoding: nonce\x00", 12)
	blk, _ := aes.NewCipher([]byte(cek))
	gcm, _ := cipher.NewGCM(blk)
	pt, err := gcm.Open(nil, []byte(nonce), ct, nil)
	if err != nil {
		return nil, err
	}
	return bytes.TrimRight(pt, "\x00")[:len(bytes.TrimRight(pt, "\x00"))-1], nil // drop the 0x02 delimiter
}

func (d *fakeDevice) payloads() []pushPayload {
	d.mu.Lock()
	defer d.mu.Unlock()
	return append([]pushPayload(nil), d.got...)
}

func waitFor(t *testing.T, what string, ok func() bool) {
	t.Helper()
	for i := 0; i < 200; i++ {
		if ok() {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", what)
}

func pushSetup(t *testing.T) (*Server, http.Handler, *User, string) {
	t.Helper()
	s, _, _ := newTestServer(t)
	s.opts.PushAllowPrivate = true
	u, err := s.createUser("Ana", "ana@example.com", "correct horse battery", "admin", "")
	if err != nil {
		t.Fatal(err)
	}
	_, key, _ := s.createAPIKey(u.ID, "test")
	return s, s.Handler(), u, key
}

func call(h http.Handler, method, path, key string, body any) *httptest.ResponseRecorder {
	var rd io.Reader
	if body != nil {
		b, _ := json.Marshal(body)
		rd = bytes.NewReader(b)
	}
	r := httptest.NewRequest(method, path, rd)
	if key != "" {
		r.Header.Set("Authorization", "Bearer "+key)
	}
	if body != nil {
		r.Header.Set("Content-Type", "application/json")
	}
	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)
	return w
}

func TestPushDeliversEncryptedMessage(t *testing.T) {
	s, h, u, key := pushSetup(t)
	dev := newFakeDevice(t)

	if w := call(h, "GET", "/api/v1/push/key", "", nil); w.Code != 401 {
		t.Fatalf("push key without sign-in: %d", w.Code)
	}
	var k struct{ PublicKey string }
	json.Unmarshal(call(h, "GET", "/api/v1/push/key", key, nil).Body.Bytes(), &k)
	if raw, err := base64.RawURLEncoding.DecodeString(k.PublicKey); err != nil || len(raw) != 65 {
		t.Fatalf("bad VAPID key %q", k.PublicKey)
	}

	w := call(h, "POST", "/api/v1/push/subscriptions", key, map[string]any{
		"endpoint": dev.srv.URL + "/push/abc", "keys": dev.keys(), "topics": []string{"alerts", "bad topic!"}, "device": "Pixel"})
	if w.Code != 200 {
		t.Fatalf("subscribe: %d %s", w.Code, w.Body)
	}
	var sub PushSub
	json.Unmarshal(w.Body.Bytes(), &sub)
	if len(sub.Topics) != 2 || sub.Topics[0] != userTopic(u.ID) || sub.Topics[1] != "alerts" {
		t.Fatalf("topics: %v", sub.Topics)
	}

	s.Publish("alerts", "Disk full", "db1 is at 98%", 5, []string{"warning"})
	s.Publish("other", "ignored", "not subscribed", 3, nil)
	waitFor(t, "push", func() bool { return len(dev.payloads()) == 1 })
	p := dev.payloads()[0]
	if p.Title != "Disk full" || p.Message != "db1 is at 98%" || p.Priority != 5 || p.Topic != "alerts" {
		t.Fatalf("payload: %+v", p)
	}

	// The VAPID header must be a valid ES256 JWT for the endpoint's origin, signed by the advertised key.
	dev.mu.Lock()
	hdr := dev.vapid
	dev.mu.Unlock()
	tok := strings.TrimPrefix(strings.Split(hdr, ", k=")[0], "vapid t=")
	parts := strings.Split(tok, ".")
	if len(parts) != 3 || !strings.HasSuffix(hdr, ", k="+k.PublicKey) {
		t.Fatalf("authorization: %s", hdr)
	}
	sig, _ := base64.RawURLEncoding.DecodeString(parts[2])
	raw, _ := base64.RawURLEncoding.DecodeString(k.PublicKey)
	pub, _ := ecdh.P256().NewPublicKey(raw)
	ek := ecdsa.PublicKey{Curve: elliptic.P256(), X: new(big.Int).SetBytes(pub.Bytes()[1:33]), Y: new(big.Int).SetBytes(pub.Bytes()[33:])}
	hh := sha256.Sum256([]byte(parts[0] + "." + parts[1]))
	if !ecdsa.Verify(&ek, hh[:], new(big.Int).SetBytes(sig[:32]), new(big.Int).SetBytes(sig[32:])) {
		t.Fatal("VAPID signature does not verify")
	}
	claims, _ := base64.RawURLEncoding.DecodeString(parts[1])
	if !strings.Contains(string(claims), `"aud":"`+dev.srv.URL+`"`) {
		t.Fatalf("claims: %s", claims)
	}

	// The key survives a restart: it is stored, not regenerated.
	s.vapid = nil
	if _, pub2, _ := s.vapidKey(); pub2 != k.PublicKey {
		t.Fatal("VAPID key changed")
	}
}

func TestPushOnCallAlertCarriesAlertID(t *testing.T) {
	s, h, u, key := pushSetup(t)
	dev := newFakeDevice(t)
	call(h, "POST", "/api/v1/push/subscriptions", key, map[string]any{"endpoint": dev.srv.URL + "/p", "keys": dev.keys(), "topics": []string{"alerts"}})
	a, _, err := s.RaiseAlert("test:cpu", "test", "CPU high", "load 12", "critical", "")
	if err != nil {
		t.Fatal(err)
	}
	waitFor(t, "push", func() bool { return len(dev.payloads()) == 1 })
	if p := dev.payloads()[0]; p.Alert != a.ID || p.Status != "firing" {
		t.Fatalf("payload: %+v", p)
	}
	if w := call(h, "POST", "/api/v1/alerts/"+a.ID+"/ack", key, map[string]any{}); w.Code != 200 {
		t.Fatalf("ack: %d %s", w.Code, w.Body)
	}
	_ = u
}

func TestPushGoneSubscriptionIsRemoved(t *testing.T) {
	s, h, _, key := pushSetup(t)
	dev := newFakeDevice(t)
	dev.status = 410
	call(h, "POST", "/api/v1/push/subscriptions", key, map[string]any{"endpoint": dev.srv.URL + "/p", "keys": dev.keys(), "topics": []string{"alerts"}})
	s.Publish("alerts", "x", "y", 3, nil)
	waitFor(t, "removal", func() bool { return len(s.pushSubs("")) == 0 })
}

func TestPushSubscriptionRules(t *testing.T) {
	_, h, _, key := pushSetup(t)
	dev := newFakeDevice(t)
	if w := call(h, "POST", "/api/v1/push/subscriptions", key, map[string]any{"endpoint": "ftp://x/y", "keys": dev.keys()}); w.Code != 400 {
		t.Fatalf("bad endpoint accepted: %d", w.Code)
	}
	if w := call(h, "POST", "/api/v1/push/subscriptions", key, map[string]any{"endpoint": dev.srv.URL, "keys": map[string]string{"p256dh": "AAAA", "auth": "AAAA"}}); w.Code != 400 {
		t.Fatalf("bad keys accepted: %d", w.Code)
	}
	body := map[string]any{"endpoint": dev.srv.URL + "/p", "keys": dev.keys(), "topics": []string{"a"}}
	call(h, "POST", "/api/v1/push/subscriptions", key, body)
	body["topics"] = []string{"b"}
	call(h, "POST", "/api/v1/push/subscriptions", key, body) // same endpoint updates in place
	var list []PushSub
	json.Unmarshal(call(h, "GET", "/api/v1/push/subscriptions", key, nil).Body.Bytes(), &list)
	if len(list) != 1 || list[0].Topics[len(list[0].Topics)-1] != "b" {
		t.Fatalf("list: %+v", list)
	}
	if w := call(h, "POST", "/api/v1/push/test", key, map[string]any{}); !strings.Contains(w.Body.String(), "Pixel") && w.Code != 200 {
		t.Fatalf("test push: %d %s", w.Code, w.Body)
	}
	if w := call(h, "DELETE", "/api/v1/push/subscriptions/"+list[0].ID, key, nil); w.Code != 200 {
		t.Fatalf("delete: %d", w.Code)
	}
}

func TestMobileLoginAndEscalate(t *testing.T) {
	s, h, _, _ := pushSetup(t)
	if w := call(h, "POST", "/api/v1/mobile/login", "", map[string]string{"email": "ana@example.com", "password": "nope nope nope"}); w.Code != 401 {
		t.Fatalf("wrong password: %d", w.Code)
	}
	w := call(h, "POST", "/api/v1/mobile/login", "", map[string]string{"email": "ana@example.com", "password": "correct horse battery", "device": "Pixel 8"})
	var out struct{ Token string }
	json.Unmarshal(w.Body.Bytes(), &out)
	if w.Code != 200 || !strings.HasPrefix(out.Token, "sb_") {
		t.Fatalf("login: %d %s", w.Code, w.Body)
	}
	if len(w.Result().Cookies()) != 0 {
		t.Fatal("mobile login must not set a cookie")
	}
	if w := call(h, "GET", "/api/v1/alerts", out.Token, nil); w.Code != 200 {
		t.Fatalf("token rejected: %d", w.Code)
	}
	if err := s.EscalateAlert("missing", "ana"); err == nil {
		t.Fatal("escalating a missing alert should fail")
	}
}

func TestAppServingAndCORS(t *testing.T) {
	s, _, _, _ := pushSetup(t)
	dir := t.TempDir()
	os.WriteFile(dir+"/index.html", []byte("<h1>app</h1>"), 0o644)
	s.opts.AppDir = dir
	s.opts.AppOrigin = "https://app.example.com"
	h := s.Handler()
	if w := call(h, "GET", "/app/", "", nil); w.Code != 200 || !strings.Contains(w.Body.String(), "<h1>app</h1>") {
		t.Fatalf("app: %d %s", w.Code, w.Body)
	}
	r := httptest.NewRequest("OPTIONS", "/api/v1/alerts", nil)
	r.Header.Set("Origin", "https://app.example.com")
	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)
	if w.Code != 204 || w.Header().Get("Access-Control-Allow-Origin") != "https://app.example.com" {
		t.Fatalf("preflight: %d %v", w.Code, w.Header())
	}
	r = httptest.NewRequest("OPTIONS", "/api/v1/alerts", nil)
	r.Header.Set("Origin", "https://evil.example.com")
	w = httptest.NewRecorder()
	h.ServeHTTP(w, r)
	if w.Header().Get("Access-Control-Allow-Origin") != "" {
		t.Fatal("foreign origin allowed")
	}
}

func TestHammerheadEdgesProxy(t *testing.T) {
	s, h, _, key := pushSetup(t)
	if w := call(h, "GET", "/api/v1/hammerhead/edges", key, nil); !strings.Contains(w.Body.String(), `"configured":false`) {
		t.Fatalf("unconfigured: %s", w.Body)
	}
	hh := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/login":
			http.SetCookie(w, &http.Cookie{Name: "hh_session", Value: "x", Path: "/"})
			io.WriteString(w, "{}")
		case "/api/edges":
			if _, err := r.Cookie("hh_session"); err != nil {
				w.WriteHeader(401)
				return
			}
			io.WriteString(w, `[{"id":"e1","name":"edge-eu","ips":["1.2.3.4"],"online":true,"token":"SECRET","tunnelAddr":"1.2.3.4:51820"}]`)
		}
	}))
	defer hh.Close()
	s.opts.HammerheadURL, s.opts.HammerheadEmail, s.opts.HammerheadPassword = hh.URL, "a@b.c", "pw"
	w := call(h, "GET", "/api/v1/hammerhead/edges", key, nil)
	if w.Code != 200 || !strings.Contains(w.Body.String(), "edge-eu") || strings.Contains(w.Body.String(), "SECRET") || strings.Contains(w.Body.String(), "51820") {
		t.Fatalf("edges: %d %s", w.Code, w.Body)
	}
}