package main

// 🤖 Sign in with Apple tests (siwa.go): a fake Apple — an httptest JWKS
// server and self-minted RS256 identity tokens — exercises the whole path
// with no Apple dependency, matching the seed-state-directly convention.

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

// fakeApple serves a JWKS for a fresh RSA key and mints identity tokens
// signed with it.
type fakeApple struct {
	key *rsa.PrivateKey
	ts  *httptest.Server
}

func newFakeApple(t *testing.T) *fakeApple {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("rsa: %v", err)
	}
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, http.StatusOK, map[string]any{
			"keys": []map[string]string{{
				"kty": "RSA", "kid": "test-kid",
				"n": base64.RawURLEncoding.EncodeToString(key.PublicKey.N.Bytes()),
				"e": base64.RawURLEncoding.EncodeToString([]byte{1, 0, 1}),
			}},
		})
	}))
	t.Cleanup(ts.Close)
	return &fakeApple{key: key, ts: ts}
}

// token mints an identity token; overrides patch the default valid claims.
func (a *fakeApple) token(t *testing.T, overrides map[string]any) string {
	t.Helper()
	claims := map[string]any{
		"iss":   appleIssuer,
		"aud":   "com.elevenmessenger.messenger",
		"sub":   "apple-sub-1",
		"exp":   time.Now().Add(5 * time.Minute).Unix(),
		"nonce": "",
	}
	for k, v := range overrides {
		claims[k] = v
	}
	head, _ := json.Marshal(map[string]string{"alg": "RS256", "kid": "test-kid"})
	body, _ := json.Marshal(claims)
	signing := base64.RawURLEncoding.EncodeToString(head) + "." + base64.RawURLEncoding.EncodeToString(body)
	digest := sha256.Sum256([]byte(signing))
	sig, err := rsa.SignPKCS1v15(rand.Reader, a.key, crypto.SHA256, digest[:])
	if err != nil {
		t.Fatalf("sign: %v", err)
	}
	return signing + "." + base64.RawURLEncoding.EncodeToString(sig)
}

// siwaServer is newTestServer with SiwA enabled against the fake Apple.
func siwaServer(t *testing.T) (*httptest.Server, *Server, *fakeApple) {
	t.Helper()
	ts, s := newTestServer(t)
	apple := newFakeApple(t)
	s.siwaClientIDs = []string{"com.elevenmessenger.messenger", "com.elevenmessenger.signin"}
	s.siwaPepper = []byte("test-pepper")
	s.siwa = &jwksCache{url: apple.ts.URL}
	return ts, s, apple
}

func siwaVerifyReq(t *testing.T, ts *httptest.Server, headers map[string]string, body map[string]any) *http.Response {
	t.Helper()
	b, _ := json.Marshal(body)
	return do(t, "POST", ts.URL+"/api/siwa/verify", "", headers, b)
}

func TestSiwaDisabledIs404(t *testing.T) {
	ts, _ := newTestServer(t) // no client ids, no pepper
	resp := siwaVerifyReq(t, ts, nil, map[string]any{"identity_token": "x"})
	if resp.StatusCode != http.StatusNotFound {
		t.Fatalf("disabled siwa: got %d, want 404", resp.StatusCode)
	}
}

func TestSiwaSignInCreatesThenResolvesAccount(t *testing.T) {
	ts, _, apple := siwaServer(t)

	// First sign-in: a fresh account, token + proof handed back.
	resp := siwaVerifyReq(t, ts, nil, map[string]any{"identity_token": apple.token(t, nil)})
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("first sign-in: got %d", resp.StatusCode)
	}
	var first struct {
		Token   string `json:"token"`
		Proof   string `json:"proof"`
		Created bool   `json:"created"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&first); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if first.Token == "" || first.Proof == "" || !first.Created {
		t.Fatalf("first sign-in: token=%q proof=%q created=%v", first.Token, first.Proof, first.Created)
	}

	// The minted token is a real device link token: it can write the vault.
	put := do(t, "PUT", ts.URL+"/api/blob", first.Token, map[string]string{"If-Match": "0"}, []byte("sealed"))
	if put.StatusCode != http.StatusOK {
		t.Fatalf("blob put with siwa token: got %d", put.StatusCode)
	}

	// Second sign-in, same sub: the SAME account (its blob is readable).
	resp2 := siwaVerifyReq(t, ts, nil, map[string]any{"identity_token": apple.token(t, nil)})
	var second struct {
		Token   string `json:"token"`
		Created bool   `json:"created"`
	}
	if err := json.NewDecoder(resp2.Body).Decode(&second); err != nil {
		t.Fatalf("decode second: %v", err)
	}
	if second.Created {
		t.Fatalf("second sign-in claims a fresh account")
	}
	get := do(t, "GET", ts.URL+"/api/blob", second.Token, nil, nil)
	if get.StatusCode != http.StatusOK {
		t.Fatalf("blob get with second token: got %d", get.StatusCode)
	}
	if b, _ := io.ReadAll(get.Body); string(b) != "sealed" {
		t.Fatalf("second sign-in reads a different vault: %q", b)
	}
}

func TestSiwaRejectsBadTokens(t *testing.T) {
	ts, _, apple := siwaServer(t)
	cases := map[string]string{
		"wrong issuer":   apple.token(t, map[string]any{"iss": "https://evil.example"}),
		"wrong audience": apple.token(t, map[string]any{"aud": "com.other.app"}),
		"expired":        apple.token(t, map[string]any{"exp": time.Now().Add(-time.Minute).Unix()}),
		"empty sub":      apple.token(t, map[string]any{"sub": ""}),
		"garbage":        "not.a.jwt",
	}
	for name, tok := range cases {
		resp := siwaVerifyReq(t, ts, nil, map[string]any{"identity_token": tok})
		if resp.StatusCode != http.StatusUnauthorized {
			t.Errorf("%s: got %d, want 401", name, resp.StatusCode)
		}
	}
	// Tampered signature: valid token, one flipped payload byte.
	tok := apple.token(t, nil)
	tampered := tok[:len(tok)-40] + "AAAA" + tok[len(tok)-36:]
	resp := siwaVerifyReq(t, ts, nil, map[string]any{"identity_token": tampered})
	if resp.StatusCode != http.StatusUnauthorized {
		t.Errorf("tampered: got %d, want 401", resp.StatusCode)
	}
}

func TestSiwaNonceBinding(t *testing.T) {
	ts, _, apple := siwaServer(t)
	tok := apple.token(t, map[string]any{"nonce": "expected-nonce"})
	resp := siwaVerifyReq(t, ts, nil, map[string]any{"identity_token": tok, "nonce": "expected-nonce"})
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("matching nonce: got %d", resp.StatusCode)
	}
	resp = siwaVerifyReq(t, ts, nil, map[string]any{"identity_token": tok, "nonce": "different"})
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("mismatched nonce: got %d, want 401", resp.StatusCode)
	}
}

func TestSiwaLinkToExistingAccount(t *testing.T) {
	ts, s, apple := siwaServer(t)
	uid, token := seedUser(t, s)
	proof := s.keys.mintCeremonyProof(uid)

	// Link needs a fresh ceremony proof — same bar as adding a passkey.
	resp := siwaVerifyReq(t, ts, nil, map[string]any{
		"identity_token": apple.token(t, map[string]any{"sub": "link-sub"}), "link": true})
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("link without proof: got %d, want 401", resp.StatusCode)
	}
	resp = siwaVerifyReq(t, ts, map[string]string{"X-Ceremony-Proof": proof}, map[string]any{
		"identity_token": apple.token(t, map[string]any{"sub": "link-sub"}), "link": true})
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("link with proof: got %d", resp.StatusCode)
	}

	// Signing in with that Apple ID now lands in the seeded account: its
	// vault is the one the seeded link token wrote.
	put := do(t, "PUT", ts.URL+"/api/blob", token, map[string]string{"If-Match": "0"}, []byte("seeded-vault"))
	if put.StatusCode != http.StatusOK {
		t.Fatalf("seed blob: got %d", put.StatusCode)
	}
	resp = siwaVerifyReq(t, ts, nil, map[string]any{
		"identity_token": apple.token(t, map[string]any{"sub": "link-sub"})})
	var signin struct {
		Token   string `json:"token"`
		Created bool   `json:"created"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&signin); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if signin.Created {
		t.Fatalf("sign-in after link created a fresh account")
	}
	get := do(t, "GET", ts.URL+"/api/blob", signin.Token, nil, nil)
	if b, _ := io.ReadAll(get.Body); string(b) != "seeded-vault" {
		t.Fatalf("linked sign-in reads a different vault: %q", b)
	}

	// The same Apple ID cannot anchor a SECOND account.
	uid2, _ := seedUser(t, s)
	proof2 := s.keys.mintCeremonyProof(uid2)
	resp = siwaVerifyReq(t, ts, map[string]string{"X-Ceremony-Proof": proof2}, map[string]any{
		"identity_token": apple.token(t, map[string]any{"sub": "link-sub"}), "link": true})
	if resp.StatusCode != http.StatusConflict {
		t.Fatalf("second-account link: got %d, want 409", resp.StatusCode)
	}
}

func TestSiwaReturnTrampoline(t *testing.T) {
	ts, _, _ := siwaServer(t)
	client := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error {
		return http.ErrUseLastResponse
	}}
	resp, err := client.PostForm(ts.URL+"/siwa/return", map[string][]string{
		"id_token": {"tok.en.value"}, "state": {"st ate"}})
	if err != nil {
		t.Fatalf("post: %v", err)
	}
	if resp.StatusCode != http.StatusSeeOther {
		t.Fatalf("trampoline: got %d, want 303", resp.StatusCode)
	}
	loc := resp.Header.Get("Location")
	want := siwaReturnScheme + "#state=st+ate&id_token=tok.en.value"
	if loc != want {
		t.Fatalf("trampoline location: got %q, want %q", loc, want)
	}
}

func TestSiwaReturnWebState(t *testing.T) {
	ts, _, _ := siwaServer(t)
	client := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error {
		return http.ErrUseLastResponse
	}}
	// A "web."-prefixed state (minted by the home page) bounces back to the
	// page itself — a relative, same-origin redirect with the token in the
	// fragment.
	resp, err := client.PostForm(ts.URL+"/siwa/return", map[string][]string{
		"id_token": {"tok.en.value"}, "state": {"web.abc123"}})
	if err != nil {
		t.Fatalf("post: %v", err)
	}
	if resp.StatusCode != http.StatusSeeOther {
		t.Fatalf("web trampoline: got %d, want 303", resp.StatusCode)
	}
	if loc, want := resp.Header.Get("Location"), "/#siwa=tok.en.value&state=web.abc123"; loc != want {
		t.Fatalf("web trampoline location: got %q, want %q", loc, want)
	}
	// Any other state keeps the native apps' custom-scheme bounce.
	resp, err = client.PostForm(ts.URL+"/siwa/return", map[string][]string{
		"id_token": {"tok.en.value"}, "state": {"native-state"}})
	if err != nil {
		t.Fatalf("post: %v", err)
	}
	if loc, want := resp.Header.Get("Location"), siwaReturnScheme+"#state=native-state&id_token=tok.en.value"; loc != want {
		t.Fatalf("native trampoline location: got %q, want %q", loc, want)
	}
}