package wristband

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

// obtainCode runs authorize, resume and (when asked) consent for userID at
// client and returns the authorization code.
func (f *oidcFixture) obtainCode(t *testing.T, client *OIDCClientRecord, userID uint, scope string) string {
	t.Helper()
	f.host.signIn(userID)
	q := oidcAuthorizeParams(f, func(q url.Values) {
		q.Set("client_id", client.ClientID)
		q.Set("redirect_uri", client.RedirectURIs[0])
		q.Set("scope", scope)
	})
	rec := f.startAndResume(t, q)
	if rec.Code == http.StatusOK {
		view := f.host.lastView(t)
		rec = f.post("/oauth/authorize", url.Values{"request_id": {view.RequestID}, "consent_token": {view.ConsentToken}, "decision": {"allow"}})
	}
	u, err := url.Parse(rec.Header().Get("Location"))
	if err != nil || u.Query().Get("code") == "" {
		t.Fatalf("no code: %d %q %s", rec.Code, rec.Header().Get("Location"), rec.Body.String())
	}
	return u.Query().Get("code")
}

// tokenCall posts form to the token endpoint, with Basic credentials when
// id is not empty.
func (f *oidcFixture) tokenCall(form url.Values, id, secret string) *httptest.ResponseRecorder {
	req := httptest.NewRequest(http.MethodPost, "/oauth/token", strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	if id != "" {
		req.SetBasicAuth(url.QueryEscape(id), url.QueryEscape(secret))
	}
	rec := httptest.NewRecorder()
	f.mux().ServeHTTP(rec, req)
	return rec
}

func (f *oidcFixture) exchange(client *OIDCClientRecord, secret, code, verifier string) *httptest.ResponseRecorder {
	return f.tokenCall(url.Values{
		"grant_type":    {"authorization_code"},
		"code":          {code},
		"redirect_uri":  {client.RedirectURIs[0]},
		"code_verifier": {verifier},
	}, client.ClientID, secret)
}

type tokenAnswer struct {
	AccessToken  string `json:"access_token"`
	IDToken      string `json:"id_token"`
	RefreshToken string `json:"refresh_token"`
	Scope        string `json:"scope"`
}

func decodeToken(t *testing.T, rec *httptest.ResponseRecorder) tokenAnswer {
	t.Helper()
	if rec.Code != http.StatusOK {
		t.Fatalf("token = %d %s", rec.Code, rec.Body.String())
	}
	var a tokenAnswer
	if err := json.Unmarshal(rec.Body.Bytes(), &a); err != nil {
		t.Fatal(err)
	}
	return a
}

func (f *oidcFixture) userinfo(token string) *httptest.ResponseRecorder {
	req := httptest.NewRequest(http.MethodGet, "/oauth/userinfo", nil)
	if token != "" {
		req.Header.Set("Authorization", "Bearer "+token)
	}
	rec := httptest.NewRecorder()
	f.mux().ServeHTTP(rec, req)
	return rec
}

func assertNoStore(t *testing.T, rec *httptest.ResponseRecorder) {
	t.Helper()
	if rec.Header().Get("Cache-Control") != "no-store, private" || rec.Header().Get("Pragma") != "no-cache" {
		t.Fatalf("headers = %v", rec.Header())
	}
}

func assertUserInfo401(t *testing.T, rec *httptest.ResponseRecorder) {
	t.Helper()
	if rec.Code != http.StatusUnauthorized || rec.Body.Len() != 0 ||
		rec.Header().Get("WWW-Authenticate") != `Bearer error="invalid_token"` ||
		rec.Header().Get("Cache-Control") != "no-store, private" {
		t.Fatalf("userinfo = %d %q %v", rec.Code, rec.Body.String(), rec.Header())
	}
}

// parseIDToken verifies an id_token against the fixture's served JWKS.
func (f *oidcFixture) parseIDToken(t *testing.T, raw string) jwt.MapClaims {
	t.Helper()
	var jwks struct {
		Keys []map[string]any `json:"keys"`
	}
	if err := json.Unmarshal(f.get("/oauth/jwks").Body.Bytes(), &jwks); err != nil {
		t.Fatal(err)
	}
	claims := jwt.MapClaims{}
	_, err := jwt.ParseWithClaims(raw, claims, func(tk *jwt.Token) (any, error) {
		for _, k := range jwks.Keys {
			if k["kid"] != tk.Header["kid"] {
				continue
			}
			switch k["kty"] {
			case "RSA":
				return rsaKeyFromJWK(t, k), nil
			case "EC":
				x, _ := base64.RawURLEncoding.DecodeString(k["x"].(string))
				y, _ := base64.RawURLEncoding.DecodeString(k["y"].(string))
				return ecdsa.ParseUncompressedPublicKey(elliptic.P256(), append(append([]byte{4}, x...), y...))
			}
		}
		return nil, fmt.Errorf("kid %v not published", tk.Header["kid"])
	}, jwt.WithValidMethods([]string{"RS256", "ES256"}), jwt.WithTimeFunc(f.srv.now))
	if err != nil {
		t.Fatalf("id_token does not verify: %v", err)
	}
	return claims
}

func TestOIDCTokenErrorsSmoke(t *testing.T) {
	f := newOIDCFixture(t, testOIDCIssuer, nil)

	// Wrong secret: 401 with the lowercase realm.
	rec := f.tokenCall(url.Values{"grant_type": {"authorization_code"}, "code": {"oas_ac_x"}, "redirect_uri": {testOIDCRedirect}, "code_verifier": {oidcTestVerifier}}, f.client.ClientID, "oas_cs_wrong")
	if rec.Code != http.StatusUnauthorized || rec.Header().Get("WWW-Authenticate") != `Basic realm="oauth"` ||
		rec.Body.String() != `{"error":"invalid_client","error_description":"Client authentication failed."}` {
		t.Fatalf("wrong secret = %d %v %s", rec.Code, rec.Header(), rec.Body.String())
	}
	assertNoStore(t, rec)

	// Unknown client and a secret without the prefix fail the same way.
	plain := f.backend.addClient(OIDCClientRecord{ClientID: "oas_plain", Kind: ClientKindOIDC, RedirectURIs: []string{testOIDCRedirect}, Scopes: []string{"openid"}, IsActive: true}, "plain-secret")
	for _, creds := range [][2]string{{"oas_unknown", testOIDCSecret}, {plain.ClientID, "plain-secret"}} {
		if rec := f.tokenCall(url.Values{"grant_type": {"authorization_code"}}, creds[0], creds[1]); rec.Code != http.StatusUnauthorized {
			t.Fatalf("%v = %d", creds, rec.Code)
		}
	}

	// Unsupported grant: 400 with no-store, private and Pragma.
	rec = f.tokenCall(url.Values{"grant_type": {"refresh_token"}}, "", "")
	if rec.Code != http.StatusBadRequest || rec.Body.String() != `{"error":"unsupported_grant_type","error_description":"Only the authorization_code grant is supported."}` {
		t.Fatalf("unsupported grant = %d %s", rec.Code, rec.Body.String())
	}
	assertNoStore(t, rec)

	// Missing parameters.
	rec = f.tokenCall(url.Values{"grant_type": {"authorization_code"}}, f.client.ClientID, f.clientSecret)
	if rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), "code, redirect_uri and code_verifier are required.") {
		t.Fatalf("missing params = %d %s", rec.Code, rec.Body.String())
	}

	// Basic credentials with percent-encoded characters authenticate.
	odd := "oas_cs_a+b/c=d~e%f"
	oddClient := f.backend.addClient(OIDCClientRecord{ClientID: "oas_odd", Kind: ClientKindOIDC, RedirectURIs: []string{testOIDCRedirect}, Scopes: []string{"openid"}, IsActive: true}, odd)
	code := f.obtainCode(t, oddClient, 7, "openid")
	decodeToken(t, f.exchange(oddClient, odd, code, oidcTestVerifier))

	// The same secret in the form works too.
	code = f.obtainCode(t, oddClient, 7, "openid")
	decodeToken(t, f.tokenCall(url.Values{
		"grant_type": {"authorization_code"}, "code": {code}, "redirect_uri": {testOIDCRedirect},
		"code_verifier": {oidcTestVerifier}, "client_id": {oddClient.ClientID}, "client_secret": {odd},
	}, "", ""))

	// An inactive client cannot exchange.
	code = f.obtainCode(t, f.client, 7, "openid")
	f.backend.clients[0].IsActive = false
	if rec := f.exchange(f.client, f.clientSecret, code, oidcTestVerifier); rec.Code != http.StatusUnauthorized {
		t.Fatalf("inactive client = %d", rec.Code)
	}
}

func TestOIDCReplayRevokesSmoke(t *testing.T) {
	f := newOIDCFixture(t, testOIDCIssuer, nil)
	code := f.obtainCode(t, f.client, 7, "openid profile")
	first := decodeToken(t, f.exchange(f.client, f.clientSecret, code, oidcTestVerifier))
	if rec := f.userinfo(first.AccessToken); rec.Code != http.StatusOK {
		t.Fatalf("userinfo before replay = %d", rec.Code)
	}

	rec := f.exchange(f.client, f.clientSecret, code, oidcTestVerifier)
	if rec.Code != http.StatusBadRequest || rec.Body.String() != `{"error":"invalid_grant","error_description":"The authorization code is invalid, expired, already used, or does not match this request."}` {
		t.Fatalf("replay = %d %s", rec.Code, rec.Body.String())
	}
	assertUserInfo401(t, f.userinfo(first.AccessToken))

	// Burn then check: a wrong verifier burns the code.
	code = f.obtainCode(t, f.client, 7, "openid profile")
	if rec := f.exchange(f.client, f.clientSecret, code, strings.Repeat("x", 43)); rec.Code != http.StatusBadRequest {
		t.Fatalf("wrong verifier = %d", rec.Code)
	}
	if rec := f.exchange(f.client, f.clientSecret, code, oidcTestVerifier); rec.Code != http.StatusBadRequest {
		t.Fatalf("burned code = %d", rec.Code)
	}

	// An expired code is invalid.
	code = f.obtainCode(t, f.client, 7, "openid profile")
	base := time.Now()
	f.srv.now = func() time.Time { return base.Add(2 * time.Minute) }
	if rec := f.exchange(f.client, f.clientSecret, code, oidcTestVerifier); rec.Code != http.StatusBadRequest {
		t.Fatalf("expired code = %d", rec.Code)
	}
}

func TestOIDCUserInfoSmoke(t *testing.T) {
	f := newOIDCFixture(t, testOIDCIssuer, nil)
	assertUserInfo401(t, f.userinfo(""))
	assertUserInfo401(t, f.userinfo("not-a-token"))
	assertUserInfo401(t, f.userinfo("oas_at_unknown"))

	// Without the email scope there are no email claims.
	tok := decodeToken(t, f.exchange(f.client, f.clientSecret, f.obtainCode(t, f.client, 7, "openid profile"), oidcTestVerifier))
	rec := f.userinfo(tok.AccessToken)
	var info map[string]any
	if err := json.Unmarshal(rec.Body.Bytes(), &info); err != nil {
		t.Fatal(err)
	}
	if info["name"] != "Ada Example" || info["email"] != nil || info["email_verified"] != nil || rec.Header().Get("Cache-Control") != "no-store, private" {
		t.Fatalf("userinfo = %s", rec.Body.String())
	}
	if f.backend.accessTokens()[0].LastUsedAt == nil {
		t.Fatal("userinfo must touch LastUsedAt")
	}

	// Revoked.
	now := time.Now()
	f.backend.tokens[0].RevokedAt = &now
	assertUserInfo401(t, f.userinfo(tok.AccessToken))

	// Expired.
	tok = decodeToken(t, f.exchange(f.client, f.clientSecret, f.obtainCode(t, f.client, 7, "openid profile"), oidcTestVerifier))
	f.srv.now = func() time.Time { return now.Add(2 * time.Hour) }
	assertUserInfo401(t, f.userinfo(tok.AccessToken))
	f.srv.now = time.Now

	// Gone user.
	delete(f.backend.identities, 7)
	assertUserInfo401(t, f.userinfo(tok.AccessToken))
}

func TestMixedKeyRingSmoke(t *testing.T) {
	f := newOIDCFixture(t, testOIDCIssuer, nil)
	before := decodeToken(t, f.exchange(f.client, f.clientSecret, f.obtainCode(t, f.client, 7, "openid"), oidcTestVerifier))

	// Rotate to ES256; the RS256 key stays published inside the grace window.
	f.keys.rotate(t, AlgES256, time.Now())
	var jwks struct {
		Keys []map[string]any `json:"keys"`
	}
	if err := json.Unmarshal(f.get("/oauth/jwks").Body.Bytes(), &jwks); err != nil {
		t.Fatal(err)
	}
	if len(jwks.Keys) != 2 || jwks.Keys[0]["kty"] != "RSA" || jwks.Keys[1]["kty"] != "EC" || jwks.Keys[1]["crv"] != "P-256" || jwks.Keys[1]["alg"] != "ES256" {
		t.Fatalf("jwks = %v", jwks.Keys)
	}
	var doc map[string]any
	if err := json.Unmarshal(f.get("/.well-known/openid-configuration").Body.Bytes(), &doc); err != nil {
		t.Fatal(err)
	}
	if fmt.Sprint(doc["id_token_signing_alg_values_supported"]) != "[RS256 ES256]" {
		t.Fatalf("algs = %v", doc["id_token_signing_alg_values_supported"])
	}

	// A token signed before rotation still verifies; a new one is ES256.
	f.parseIDToken(t, before.IDToken)
	after := decodeToken(t, f.exchange(f.client, f.clientSecret, f.obtainCode(t, f.client, 7, "openid"), oidcTestVerifier))
	tok, _, err := jwt.NewParser().ParseUnverified(after.IDToken, jwt.MapClaims{})
	if err != nil || tok.Header["alg"] != "ES256" {
		t.Fatalf("new id_token alg = %v (%v)", tok.Header["alg"], err)
	}
	f.parseIDToken(t, after.IDToken)

	// After the grace window only the EC key is published.
	f.srv.now = func() time.Time { return time.Now().Add(25 * time.Hour) }
	if err := json.Unmarshal(f.get("/oauth/jwks").Body.Bytes(), &jwks); err != nil {
		t.Fatal(err)
	}
	if len(jwks.Keys) != 1 || jwks.Keys[0]["kty"] != "EC" {
		t.Fatalf("jwks after grace = %v", jwks.Keys)
	}

	// The key helpers refuse what they do not support.
	if _, _, _, err := GenerateSigningKey("HS256"); err == nil {
		t.Fatal("GenerateSigningKey must refuse HS256")
	}
	_, ecPriv, _, _ := GenerateSigningKey(AlgES256)
	signer, _ := ParsePrivateKeyPEM(ecPriv)
	if _, err := signIDToken(SigningKey{Kid: "k", Alg: AlgRS256, Signer: signer}, idTokenInput{}); err == nil {
		t.Fatal("an EC key labelled RS256 must be refused")
	}
}

func TestPairwiseSubjectSmoke(t *testing.T) {
	f := newOIDCFixture(t, testOIDCIssuer, nil)
	a := f.backend.addClient(OIDCClientRecord{ClientID: "oas_a", Kind: ClientKindOIDC, RedirectURIs: []string{"https://a.example.test/cb"}, Scopes: []string{"openid"}, IsActive: true, SubjectType: SubjectTypePairwise}, testOIDCSecret)
	b := f.backend.addClient(OIDCClientRecord{ClientID: "oas_b", Kind: ClientKindOIDC, RedirectURIs: []string{"https://b.example.test/cb"}, Scopes: []string{"openid"}, IsActive: true, SubjectType: SubjectTypePairwise}, testOIDCSecret)

	subOf := func(c *OIDCClientRecord) (string, string) {
		tok := decodeToken(t, f.exchange(c, testOIDCSecret, f.obtainCode(t, c, 7, "openid"), oidcTestVerifier))
		claims := f.parseIDToken(t, tok.IDToken)
		var info map[string]any
		if err := json.Unmarshal(f.userinfo(tok.AccessToken).Body.Bytes(), &info); err != nil {
			t.Fatal(err)
		}
		return claims["sub"].(string), info["sub"].(string)
	}
	subA, infoA := subOf(a)
	subB, infoB := subOf(b)
	subA2, _ := subOf(a)
	subPublic, _ := subOf(f.client)
	if subA == subB || subA == subPublic || subB == subPublic {
		t.Fatalf("pairwise subs must differ: a=%s b=%s public=%s", subA, subB, subPublic)
	}
	if subA != subA2 || subA != infoA || subB != infoB {
		t.Fatalf("subs must be stable and match userinfo: %s %s %s / %s %s", subA, subA2, infoA, subB, infoB)
	}

	mixed := &OIDCClientRecord{SubjectType: SubjectTypePairwise, RedirectURIs: []string{"https://a.example.test/cb", "https://b.example.test/cb"}}
	if _, err := SubjectSector(mixed); !errors.Is(err, ErrSectorRequired) {
		t.Fatalf("mixed hosts = %v", err)
	}
	mixed.SectorIdentifier = " sector.example.test "
	if s, err := SubjectSector(mixed); err != nil || s != "sector.example.test" {
		t.Fatalf("sector = %q %v", s, err)
	}
	if s, _ := SubjectSector(&OIDCClientRecord{SubjectType: SubjectTypePairwise, RedirectURIs: []string{"https://A.example.test/x", "https://a.example.test/y"}}); s != "a.example.test" {
		t.Fatalf("single host sector = %q", s)
	}
}

func TestOIDCRefreshOptInSmoke(t *testing.T) {
	enabled := newOIDCFixture(t, testOIDCIssuer, func(o *OIDCOptions) { o.RefreshEnabled = true })
	offline := enabled.backend.addClient(OIDCClientRecord{
		ClientID: "oas_offline", Kind: ClientKindOIDC, RedirectURIs: []string{testOIDCRedirect},
		Scopes: []string{"openid", "profile", "offline_access"}, IsActive: true, OfflineAccess: true,
	}, testOIDCSecret)

	var doc map[string]any
	if err := json.Unmarshal(enabled.get("/.well-known/openid-configuration").Body.Bytes(), &doc); err != nil {
		t.Fatal(err)
	}
	if fmt.Sprint(doc["grant_types_supported"]) != "[authorization_code refresh_token]" || !strings.Contains(fmt.Sprint(doc["scopes_supported"]), "offline_access") {
		t.Fatalf("discovery = %v %v", doc["grant_types_supported"], doc["scopes_supported"])
	}

	// A client without OfflineAccess cannot ask for offline_access, gets no
	// refresh token and may not use the grant.
	rec := enabled.get("/oauth/authorize?" + oidcAuthorizeParams(enabled, func(q url.Values) { q.Set("scope", "openid offline_access") }).Encode())
	if !strings.Contains(rec.Header().Get("Location"), "not+allowed+for+this+client") {
		t.Fatalf("offline_access without OfflineAccess = %q", rec.Header().Get("Location"))
	}
	plain := decodeToken(t, enabled.exchange(enabled.client, enabled.clientSecret, enabled.obtainCode(t, enabled.client, 7, "openid"), oidcTestVerifier))
	if plain.RefreshToken != "" {
		t.Fatal("a client without OfflineAccess must not get a refresh token")
	}
	rec = enabled.tokenCall(url.Values{"grant_type": {"refresh_token"}, "refresh_token": {"oas_rt_x"}}, enabled.client.ClientID, enabled.clientSecret)
	if rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), `"unauthorized_client"`) {
		t.Fatalf("refresh without OfflineAccess = %d %s", rec.Code, rec.Body.String())
	}

	// OfflineAccess client granted offline_access: rotation works.
	first := decodeToken(t, enabled.exchange(offline, testOIDCSecret, enabled.obtainCode(t, offline, 7, "openid profile offline_access"), oidcTestVerifier))
	if !strings.HasPrefix(first.RefreshToken, "oas_rt_") {
		t.Fatalf("refresh token = %q", first.RefreshToken)
	}
	refresh := func(token, scope string) *httptest.ResponseRecorder {
		form := url.Values{"grant_type": {"refresh_token"}, "refresh_token": {token}}
		if scope != "" {
			form.Set("scope", scope)
		}
		return enabled.tokenCall(form, offline.ClientID, testOIDCSecret)
	}
	if rec := refresh(first.RefreshToken, "openid email"); rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), `"invalid_scope"`) {
		t.Fatalf("scope beyond grant = %d %s", rec.Code, rec.Body.String())
	}
	second := decodeToken(t, refresh(first.RefreshToken, "openid profile"))
	if second.RefreshToken == "" || second.RefreshToken == first.RefreshToken || second.Scope != "openid profile" {
		t.Fatalf("rotation = %+v", second)
	}
	claims := enabled.parseIDToken(t, second.IDToken)
	if _, ok := claims["nonce"]; ok || claims["at_hash"] == nil {
		t.Fatalf("refreshed id_token claims = %v", claims)
	}
	if rec := enabled.userinfo(second.AccessToken); rec.Code != http.StatusOK {
		t.Fatalf("userinfo with refreshed token = %d", rec.Code)
	}

	// Replaying the rotated token revokes the lineage and its access tokens.
	if rec := refresh(first.RefreshToken, ""); rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), `"invalid_grant"`) {
		t.Fatalf("replay = %d %s", rec.Code, rec.Body.String())
	}
	if rec := refresh(second.RefreshToken, ""); rec.Code != http.StatusBadRequest {
		t.Fatalf("successor after replay = %d", rec.Code)
	}
	assertUserInfo401(t, enabled.userinfo(second.AccessToken))

	// A disabled server answers the grant exactly like the PHP provider.
	disabled := newOIDCFixture(t, testOIDCIssuer, nil)
	rec = disabled.tokenCall(url.Values{"grant_type": {"refresh_token"}, "refresh_token": {first.RefreshToken}}, disabled.client.ClientID, disabled.clientSecret)
	if rec.Body.String() != `{"error":"unsupported_grant_type","error_description":"Only the authorization_code grant is supported."}` {
		t.Fatalf("disabled = %s", rec.Body.String())
	}

	// Credentials for a new client follow the PHP shapes.
	id, secret, hash, err := IssueOIDCClientCredentials(DefaultOIDCOptions())
	if err != nil || len(id) != len("oas_")+24 || !strings.HasPrefix(id, "oas_") || !strings.HasPrefix(secret, "oas_cs_") || hash != sha256Hex(secret) {
		t.Fatalf("credentials = %q %q %q %v", id, secret, hash, err)
	}
}