package auth

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	jose "github.com/go-jose/go-jose/v4"

	"gitea.parodia.dev/drwily/petal/internal/db"
)

// These tests run the whole login round-trip against a stub identity provider:
// discovery, the redirect out, the callback back, and the session that comes out
// the other end. The flow is the one place in Petal where getting a detail wrong
// (an unchecked state, a nonce nobody compares) is both easy and invisible —
// everything still "works" from the browser's point of view.

// stubIdP is a minimal OpenID provider: discovery, a JWKS, and a token endpoint
// that mints a signed ID token for whoever the test says just logged in.
type stubIdP struct {
	*httptest.Server
	key      *rsa.PrivateKey
	clientID string
	// issuer as advertised by discovery and asserted in tokens. Defaults to the
	// server's URL; a test can give it a trailing slash, which is what Authentik
	// does and which OIDC requires to match byte-for-byte.
	issuer string

	// Claims the next token exchange will assert.
	sub, email, name string
	// nonce echoed into the token; set from the login attempt's cookie.
	nonce string
	// lastForm records what Petal sent to /token, so the test can assert PKCE.
	lastForm url.Values
}

func newStubIdP(t *testing.T, clientID string) *stubIdP {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	idp := &stubIdP{key: key, clientID: clientID}

	mux := http.NewServeMux()
	idp.Server = httptest.NewServer(mux)
	idp.issuer = idp.URL
	t.Cleanup(idp.Close)

	mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]any{
			"issuer":                                idp.issuer,
			"authorization_endpoint":                idp.URL + "/authorize",
			"token_endpoint":                        idp.URL + "/token",
			"jwks_uri":                              idp.URL + "/jwks",
			"id_token_signing_alg_values_supported": []string{"RS256"},
		})
	})

	mux.HandleFunc("/jwks", func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(jose.JSONWebKeySet{
			Keys: []jose.JSONWebKey{{Key: key.Public(), Algorithm: "RS256", Use: "sig"}},
		})
	})

	mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) {
		_ = r.ParseForm()
		idp.lastForm = r.PostForm
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{
			"access_token": "stub-access-token",
			"token_type":   "Bearer",
			"id_token":     idp.idToken(t),
		})
	})

	return idp
}

// idToken mints a signed ID token asserting the currently configured claims.
func (idp *stubIdP) idToken(t *testing.T) string {
	t.Helper()
	signer, err := jose.NewSigner(
		jose.SigningKey{Algorithm: jose.RS256, Key: idp.key},
		(&jose.SignerOptions{}).WithType("JWT"),
	)
	if err != nil {
		t.Fatal(err)
	}
	payload, _ := json.Marshal(map[string]any{
		"iss":   idp.issuer,
		"aud":   idp.clientID,
		"sub":   idp.sub,
		"email": idp.email,
		"name":  idp.name,
		"nonce": idp.nonce,
		"exp":   time.Now().Add(time.Hour).Unix(),
		"iat":   time.Now().Unix(),
	})
	signed, err := signer.Sign(payload)
	if err != nil {
		t.Fatal(err)
	}
	raw, err := signed.CompactSerialize()
	if err != nil {
		t.Fatal(err)
	}
	return raw
}

// newFlow wires Petal's login routes to a stub provider.
func newFlow(t *testing.T, allowed Allowlist) (*stubIdP, http.Handler, *SessionStore, *UserStore) {
	t.Helper()
	sessions, users, _ := newStores(t)
	idp := newStubIdP(t, "petal")

	o := NewOIDC(context.Background(), Options{
		IssuerURL:    idp.URL,
		ClientID:     "petal",
		ClientSecret: "shh",
		BaseURL:      "http://petal.test",
		Allowed:      allowed,
	}, sessions, users)
	return idp, o.Routes(), sessions, users
}

// cookieJar collects Set-Cookie headers across the redirect chain, standing in
// for the browser that would normally carry them.
type cookieJar map[string]string

func (j cookieJar) absorb(rec *httptest.ResponseRecorder) {
	for _, c := range rec.Result().Cookies() {
		if c.MaxAge < 0 || c.Value == "" {
			delete(j, c.Name)
			continue
		}
		j[c.Name] = c.Value
	}
}

func (j cookieJar) attach(r *http.Request) *http.Request {
	for name, value := range j {
		r.AddCookie(&http.Cookie{Name: name, Value: value})
	}
	return r
}

// start runs /auth/login and returns the redirect target plus the cookies it set.
func start(t *testing.T, flow http.Handler) (*url.URL, cookieJar) {
	t.Helper()
	rec := httptest.NewRecorder()
	flow.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/login", nil))
	if rec.Code != http.StatusFound {
		t.Fatalf("login status=%d body=%s", rec.Code, rec.Body)
	}
	target, err := url.Parse(rec.Header().Get("Location"))
	if err != nil {
		t.Fatal(err)
	}
	jar := cookieJar{}
	jar.absorb(rec)
	return target, jar
}

func TestLoginRoundTrip(t *testing.T) {
	idp, flow, sessions, users := newFlow(t, nil)
	idp.sub, idp.email, idp.name = "sub-her", "her@example.com", "Her Name"

	target, jar := start(t, flow)

	// The redirect must carry everything the flow depends on later.
	q := target.Query()
	if q.Get("state") == "" || q.Get("nonce") == "" {
		t.Fatalf("login redirect missing state/nonce: %s", target)
	}
	if q.Get("code_challenge") == "" || q.Get("code_challenge_method") != "S256" {
		t.Fatalf("login redirect missing PKCE challenge: %s", target)
	}
	if q.Get("redirect_uri") != "http://petal.test/auth/callback" {
		t.Fatalf("redirect_uri = %q", q.Get("redirect_uri"))
	}
	if jar[stateCookie] != q.Get("state") {
		t.Fatal("the state cookie does not match the state sent to the provider")
	}
	idp.nonce = jar[nonceCookie]

	// Come back as the provider would.
	rec := httptest.NewRecorder()
	flow.ServeHTTP(rec, jar.attach(
		httptest.NewRequest(http.MethodGet, "/callback?code=abc&state="+url.QueryEscape(q.Get("state")), nil)))
	if rec.Code != http.StatusFound || rec.Header().Get("Location") != "/" {
		t.Fatalf("callback status=%d location=%q body=%s", rec.Code, rec.Header().Get("Location"), rec.Body)
	}

	// PKCE: the code verifier must reach the token endpoint.
	if v := idp.lastForm.Get("code_verifier"); v == "" {
		t.Fatal("token exchange sent no code_verifier")
	}

	// The account was provisioned from the token's claims...
	user, err := users.Get("sub-her")
	if err != nil {
		t.Fatalf("user was not provisioned: %v", err)
	}
	if user.Email != "her@example.com" || user.DisplayName != "Her Name" {
		t.Fatalf("unexpected provisioned user %+v", user)
	}

	// ...and the response carries a session that resolves to them.
	jar.absorb(rec)
	token := jar[SessionCookie]
	if token == "" {
		t.Fatal("callback issued no session cookie")
	}
	got, err := sessions.Resolve(withCookie(token))
	if err != nil || got != "sub-her" {
		t.Fatalf("session resolved to %q (err=%v), want sub-her", got, err)
	}

	// The one-shot login cookies must not linger.
	for _, name := range []string{stateCookie, nonceCookie, pkceCookie} {
		if jar[name] != "" {
			t.Fatalf("%s survived the callback", name)
		}
	}

	// Signing out revokes server-side, not just in the browser.
	out := httptest.NewRecorder()
	flow.ServeHTTP(out, jar.attach(httptest.NewRequest(http.MethodPost, "/logout", nil)))
	if out.Code != http.StatusFound {
		t.Fatalf("logout status=%d", out.Code)
	}
	if _, err := sessions.Resolve(withCookie(token)); err == nil {
		t.Fatal("the session survived signing out")
	}
}

// Signing out is a state change, so it must not be reachable by GET: with
// SameSite=Lax the session cookie *is* sent on a top-level cross-site
// navigation, which would let any page on the internet sign her out mid-draft.
func TestLogoutRejectsGET(t *testing.T) {
	_, flow, _, _ := newFlow(t, nil)

	out := httptest.NewRecorder()
	flow.ServeHTTP(out, httptest.NewRequest(http.MethodGet, "/logout", nil))
	if out.Code != http.StatusMethodNotAllowed {
		t.Fatalf("GET /logout status=%d, want 405", out.Code)
	}
}

// A callback whose state doesn't match the cookie is a forged one.
func TestCallbackRejectsBadState(t *testing.T) {
	idp, flow, sessions, _ := newFlow(t, nil)
	idp.sub, idp.email = "sub-her", "her@example.com"

	_, jar := start(t, flow)
	idp.nonce = jar[nonceCookie]

	rec := httptest.NewRecorder()
	flow.ServeHTTP(rec, jar.attach(
		httptest.NewRequest(http.MethodGet, "/callback?code=abc&state=some-other-state", nil)))
	if rec.Code != http.StatusBadRequest {
		t.Fatalf("status=%d, want 400", rec.Code)
	}
	assertNoSession(t, sessions, rec)

	// And so is one with no state cookie at all.
	bare := httptest.NewRecorder()
	flow.ServeHTTP(bare, httptest.NewRequest(http.MethodGet, "/callback?code=abc&state=x", nil))
	if bare.Code != http.StatusBadRequest {
		t.Fatalf("status=%d for a cookieless callback, want 400", bare.Code)
	}
}

// An ID token minted for a different login attempt must not be accepted, even
// though it is perfectly valid and correctly signed.
func TestCallbackRejectsReplayedNonce(t *testing.T) {
	idp, flow, sessions, _ := newFlow(t, nil)
	idp.sub, idp.email = "sub-her", "her@example.com"

	_, jarA := start(t, flow)
	_, jarB := start(t, flow)
	idp.nonce = jarB[nonceCookie] // a token belonging to the *other* attempt

	rec := httptest.NewRecorder()
	flow.ServeHTTP(rec, jarA.attach(
		httptest.NewRequest(http.MethodGet, "/callback?code=abc&state="+url.QueryEscape(jarA[stateCookie]), nil)))
	if rec.Code != http.StatusBadRequest {
		t.Fatalf("status=%d, want 400", rec.Code)
	}
	assertNoSession(t, sessions, rec)
}

// Being a valid user at the identity provider is not the same as being a user
// here, and the refusal has to read like Petal rather than like a stack trace.
func TestCallbackHonoursAllowlist(t *testing.T) {
	idp, flow, sessions, users := newFlow(t, ParseAllowlist("her@example.com"))
	idp.sub, idp.email, idp.name = "sub-stranger", "stranger@example.com", "A Stranger"

	_, jar := start(t, flow)
	idp.nonce = jar[nonceCookie]

	rec := httptest.NewRecorder()
	flow.ServeHTTP(rec, jar.attach(
		httptest.NewRequest(http.MethodGet, "/callback?code=abc&state="+url.QueryEscape(jar[stateCookie]), nil)))
	if rec.Code != http.StatusForbidden {
		t.Fatalf("status=%d, want 403", rec.Code)
	}
	if body := rec.Body.String(); !strings.Contains(body, "这个 Petal 不是给你写的") ||
		!strings.Contains(body, "isn't yours to write in") {
		t.Fatalf("refusal page is not the warm bilingual one: %s", body)
	}
	assertNoSession(t, sessions, rec)
	if _, err := users.Get("sub-stranger"); err == nil {
		t.Fatal("a rejected login still provisioned an account")
	}

	// The person on the list gets in through the same door.
	idp.sub, idp.email, idp.name = "sub-her", "her@example.com", "Her Name"
	_, jar2 := start(t, flow)
	idp.nonce = jar2[nonceCookie]
	ok := httptest.NewRecorder()
	flow.ServeHTTP(ok, jar2.attach(
		httptest.NewRequest(http.MethodGet, "/callback?code=abc&state="+url.QueryEscape(jar2[stateCookie]), nil)))
	if ok.Code != http.StatusFound {
		t.Fatalf("an allowed writer was turned away: status=%d body=%s", ok.Code, ok.Body)
	}
}

// The provider refusing the login (a cancelled consent, a locked account) is a
// dead end, not a session.
func TestCallbackHandlesProviderError(t *testing.T) {
	_, flow, sessions, _ := newFlow(t, nil)
	rec := httptest.NewRecorder()
	flow.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/callback?error=access_denied", nil))
	if rec.Code != http.StatusForbidden {
		t.Fatalf("status=%d, want 403", rec.Code)
	}
	assertNoSession(t, sessions, rec)
}

// Authentik's issuer ends in a slash, and OIDC requires the discovered issuer to
// match the configured one byte-for-byte. Normalising it away made discovery
// fail against the real provider while every stub test still passed.
func TestDiscoveryKeepsTrailingSlashIssuer(t *testing.T) {
	sessions, users, _ := newStores(t)
	idp := newStubIdP(t, "petal")
	idp.issuer = idp.URL + "/"
	idp.sub, idp.email = "sub-her", "her@example.com"

	o := NewOIDC(context.Background(), Options{
		IssuerURL:    idp.issuer,
		ClientID:     "petal",
		ClientSecret: "shh",
		BaseURL:      "http://petal.test",
	}, sessions, users)
	flow := o.Routes()

	// A failed discovery renders the 503 "sign-in is unavailable" page instead
	// of redirecting, so reaching the provider at all is the assertion.
	target, jar := start(t, flow)
	if !strings.HasPrefix(target.String(), idp.URL+"/authorize") {
		t.Fatalf("login went to %q, want the provider's authorize endpoint", target)
	}

	// And the ID token it issues, whose `iss` carries the same slash, verifies.
	idp.nonce = jar[nonceCookie]
	rec := httptest.NewRecorder()
	flow.ServeHTTP(rec, jar.attach(
		httptest.NewRequest(http.MethodGet, "/callback?code=abc&state="+url.QueryEscape(jar[stateCookie]), nil)))
	if rec.Code != http.StatusFound {
		t.Fatalf("callback status=%d body=%s", rec.Code, rec.Body)
	}
}

// Signing in when already signed in shouldn't bounce a good session through the
// identity provider.
func TestLoginSkipsWhenAlreadySignedIn(t *testing.T) {
	_, flow, sessions, _ := newFlow(t, nil)
	token, err := sessions.Create(db.LocalUserID, "")
	if err != nil {
		t.Fatal(err)
	}
	rec := httptest.NewRecorder()
	flow.ServeHTTP(rec, withCookie(token))
	// withCookie builds a GET "/" request; point it at the login route.
	rec = httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/login", nil)
	req.AddCookie(&http.Cookie{Name: SessionCookie, Value: token})
	flow.ServeHTTP(rec, req)

	if rec.Code != http.StatusFound || rec.Header().Get("Location") != "/" {
		t.Fatalf("status=%d location=%q, want a redirect home", rec.Code, rec.Header().Get("Location"))
	}
}

// assertNoSession fails if a response handed out a usable session cookie.
func assertNoSession(t *testing.T, sessions *SessionStore, rec *httptest.ResponseRecorder) {
	t.Helper()
	for _, c := range rec.Result().Cookies() {
		if c.Name == SessionCookie && c.Value != "" {
			if _, err := sessions.Resolve(withCookie(c.Value)); err == nil {
				t.Fatal("a rejected login was given a working session")
			}
		}
	}
}