package wristband

import (
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"math/big"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"

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

// oidcTestVerifier is a fixed PKCE verifier (43 characters).
const oidcTestVerifier = "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"

// noRedirectClient returns an HTTP client that never follows redirects.
func noRedirectClient() *http.Client {
	return &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
}

func readBody(t *testing.T, res *http.Response) []byte {
	t.Helper()
	defer res.Body.Close()
	b, err := io.ReadAll(res.Body)
	if err != nil {
		t.Fatal(err)
	}
	return b
}

// startOIDCTestServer serves the fixture's routes on an httptest server
// whose URL is the issuer.
func startOIDCTestServer(t *testing.T, mutate func(*OIDCOptions)) (*oidcFixture, *httptest.Server) {
	t.Helper()
	ts := httptest.NewUnstartedServer(nil)
	f := newOIDCFixture(t, "http://"+ts.Listener.Addr().String(), mutate)
	ts.Config.Handler = f.mux()
	ts.Start()
	t.Cleanup(ts.Close)
	return f, ts
}

// rsaKeyFromJWK rebuilds an RSA public key from a served JWK with the
// standard library only.
func rsaKeyFromJWK(t *testing.T, k map[string]any) *rsa.PublicKey {
	t.Helper()
	n, err := base64.RawURLEncoding.DecodeString(k["n"].(string))
	if err != nil {
		t.Fatal(err)
	}
	e, err := base64.RawURLEncoding.DecodeString(k["e"].(string))
	if err != nil {
		t.Fatal(err)
	}
	return &rsa.PublicKey{N: new(big.Int).SetBytes(n), E: int(new(big.Int).SetBytes(e).Int64())}
}

// TestOIDCTracerSmoke drives a relying party end to end over the memory
// fakes: discovery, JWKS, authorize as a signed-in user, resume, consent,
// code exchange with client_secret_basic, id_token verification with
// golang-jwt against the served JWKS, and userinfo.
func TestOIDCTracerSmoke(t *testing.T) {
	f, ts := startOIDCTestServer(t, nil)
	hc := noRedirectClient()

	// Discovery: the relying party uses only the URLs it reads here.
	res, err := hc.Get(ts.URL + "/.well-known/openid-configuration")
	if err != nil {
		t.Fatal(err)
	}
	if res.StatusCode != http.StatusOK {
		t.Fatalf("discovery status = %d", res.StatusCode)
	}
	if cc := res.Header.Get("Cache-Control"); cc != "max-age=3600, public" {
		t.Fatalf("discovery Cache-Control = %q", cc)
	}
	var doc map[string]any
	if err := json.Unmarshal(readBody(t, res), &doc); err != nil {
		t.Fatal(err)
	}
	if doc["issuer"] != ts.URL {
		t.Fatalf("issuer = %v, want %s", doc["issuer"], ts.URL)
	}
	if fmt.Sprint(doc["id_token_signing_alg_values_supported"]) != "[RS256]" {
		t.Fatalf("algs = %v", doc["id_token_signing_alg_values_supported"])
	}

	// JWKS.
	res, err = hc.Get(doc["jwks_uri"].(string))
	if err != nil {
		t.Fatal(err)
	}
	var jwks struct {
		Keys []map[string]any `json:"keys"`
	}
	if err := json.Unmarshal(readBody(t, res), &jwks); err != nil {
		t.Fatal(err)
	}
	if len(jwks.Keys) != 1 || jwks.Keys[0]["kty"] != "RSA" || jwks.Keys[0]["use"] != "sig" || jwks.Keys[0]["alg"] != "RS256" {
		t.Fatalf("jwks = %v", jwks.Keys)
	}

	// Authorize as a signed-in user.
	f.host.signIn(7)
	authz := doc["authorization_endpoint"].(string)
	q := url.Values{
		"response_type":         {"code"},
		"client_id":             {f.client.ClientID},
		"redirect_uri":          {testOIDCRedirect},
		"scope":                 {"openid profile email"},
		"state":                 {"st-1"},
		"nonce":                 {"n-123"},
		"code_challenge":        {s256Challenge(oidcTestVerifier)},
		"code_challenge_method": {"S256"},
	}
	res, err = hc.Get(authz + "?" + q.Encode())
	if err != nil {
		t.Fatal(err)
	}
	readBody(t, res)
	loc := res.Header.Get("Location")
	if res.StatusCode != http.StatusFound || !strings.HasPrefix(loc, "/oauth/authorize/resume/") {
		t.Fatalf("authorize = %d %q", res.StatusCode, loc)
	}

	// Resume renders consent through the host.
	res, err = hc.Get(ts.URL + loc)
	if err != nil {
		t.Fatal(err)
	}
	readBody(t, res)
	if res.StatusCode != http.StatusOK || res.Header.Get("X-Frame-Options") != "DENY" {
		t.Fatalf("resume = %d, X-Frame-Options %q", res.StatusCode, res.Header.Get("X-Frame-Options"))
	}
	view := f.host.lastView(t)
	if strings.Join(view.Scopes, " ") != "profile email" || view.ConsentToken == "" {
		t.Fatalf("view = %+v", view)
	}

	// Allow.
	res, err = hc.PostForm(authz, url.Values{
		"request_id":    {view.RequestID},
		"consent_token": {view.ConsentToken},
		"decision":      {"allow"},
	})
	if err != nil {
		t.Fatal(err)
	}
	readBody(t, res)
	back, err := url.Parse(res.Header.Get("Location"))
	if err != nil || res.StatusCode != http.StatusFound || !strings.HasPrefix(back.String(), testOIDCRedirect+"?code=oas_ac_") {
		t.Fatalf("decide = %d %q", res.StatusCode, res.Header.Get("Location"))
	}
	if back.Query().Get("state") != "st-1" || back.Query().Has("iss") {
		t.Fatalf("decide query = %v", back.Query())
	}
	code := back.Query().Get("code")

	// Exchange with client_secret_basic, credentials form-encoded as the
	// x/oauth2 client sends them.
	form := url.Values{
		"grant_type":    {"authorization_code"},
		"code":          {code},
		"redirect_uri":  {testOIDCRedirect},
		"code_verifier": {oidcTestVerifier},
	}
	req, _ := http.NewRequest(http.MethodPost, doc["token_endpoint"].(string), strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.SetBasicAuth(url.QueryEscape(f.client.ClientID), url.QueryEscape(f.clientSecret))
	res, err = hc.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	body := readBody(t, res)
	if res.StatusCode != http.StatusOK {
		t.Fatalf("token = %d %s", res.StatusCode, body)
	}
	if res.Header.Get("Cache-Control") != "no-store, private" || res.Header.Get("Pragma") != "no-cache" {
		t.Fatalf("token headers = %v", res.Header)
	}
	var tok struct {
		AccessToken string `json:"access_token"`
		TokenType   string `json:"token_type"`
		ExpiresIn   int64  `json:"expires_in"`
		IDToken     string `json:"id_token"`
		Scope       string `json:"scope"`
	}
	if err := json.Unmarshal(body, &tok); err != nil {
		t.Fatal(err)
	}
	if tok.TokenType != "Bearer" || tok.ExpiresIn != 3600 || tok.Scope != "openid profile email" || !strings.HasPrefix(tok.AccessToken, "oas_at_") {
		t.Fatalf("token body = %s", body)
	}

	// Verify the id_token with golang-jwt against the served JWKS only.
	keyfunc := func(tk *jwt.Token) (any, error) {
		kid, _ := tk.Header["kid"].(string)
		for _, k := range jwks.Keys {
			if k["kid"] == kid {
				return rsaKeyFromJWK(t, k), nil
			}
		}
		return nil, fmt.Errorf("unknown kid %q", kid)
	}
	claims := jwt.MapClaims{}
	if _, err := jwt.ParseWithClaims(tok.IDToken, claims, keyfunc,
		jwt.WithValidMethods([]string{"RS256"}), jwt.WithIssuer(ts.URL), jwt.WithAudience(f.client.ClientID)); err != nil {
		t.Fatalf("id_token does not verify: %v", err)
	}
	payload, err := base64.RawURLEncoding.DecodeString(strings.Split(tok.IDToken, ".")[1])
	if err != nil {
		t.Fatal(err)
	}
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(payload, &raw); err != nil {
		t.Fatal(err)
	}
	if string(raw["aud"]) != `"`+f.client.ClientID+`"` {
		t.Fatalf("aud = %s, want a JSON string", raw["aud"])
	}
	wantSub := f.backend.subject(7, "")
	if wantSub == "" || claims["sub"] != wantSub {
		t.Fatalf("sub = %v, want %q", claims["sub"], wantSub)
	}
	if claims["nonce"] != "n-123" {
		t.Fatalf("nonce = %v", claims["nonce"])
	}
	sum := sha256.Sum256([]byte(tok.AccessToken))
	if claims["at_hash"] != base64.RawURLEncoding.EncodeToString(sum[:16]) {
		t.Fatalf("at_hash = %v", claims["at_hash"])
	}
	if claims["iss"] != ts.URL || claims["email"] != "ada@example.test" || claims["x_extra"] != nil {
		t.Fatalf("claims = %v", claims)
	}
	if _, ok := claims["auth_time"]; ok {
		t.Fatal("auth_time must not be emitted")
	}
	if _, ok := claims["azp"]; ok {
		t.Fatal("azp must not be emitted")
	}

	// Userinfo answers the same subject.
	req, _ = http.NewRequest(http.MethodGet, doc["userinfo_endpoint"].(string), nil)
	req.Header.Set("Authorization", "Bearer "+tok.AccessToken)
	res, err = hc.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	body = readBody(t, res)
	if res.StatusCode != http.StatusOK {
		t.Fatalf("userinfo = %d %s", res.StatusCode, body)
	}
	var info map[string]any
	if err := json.Unmarshal(body, &info); err != nil {
		t.Fatal(err)
	}
	if info["sub"] != wantSub || info["email"] != "ada@example.test" || info["iss"] != nil {
		t.Fatalf("userinfo = %s", body)
	}
	if !strings.HasPrefix(string(body), `{"sub":`) {
		t.Fatalf("userinfo must start with sub: %s", body)
	}
}