// Package oidcstub provides a minimal in-process OIDC provider for testing.
// It serves discovery, JWKS, and a /dev/token convenience endpoint from a
// single http.Handler that callers can mount on any httptest.Server.
//
// Usage:
//
//	stub, _ := oidcstub.New(oidcstub.Options{Issuer: "http://placeholder"})
//	srv := httptest.NewServer(stub.Handler())
//	stub.SetIssuer(srv.URL)
//	tok, _ := stub.SignToken(map[string]any{"sub": "alice", "aud": "myapp"})
//
// IV1: imports only stdlib + go.bigb.es/auxilia/culpa.
package oidcstub

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"math/big"
	"net/http"
	"net/url"
	"time"

	"go.bigb.es/auxilia/culpa"
)

const defaultDevStubUser = "bigbes"

const defaultKID = "oidcstub-key-1"

// Options configures a Stub.
type Options struct {
	// Issuer is the OIDC issuer URL. Required.
	Issuer string
	// Audience is the default audience claim included by Mint. Optional.
	Audience string
	// UsernameClaim is the JWT claim set to the sub argument by Mint. Optional.
	UsernameClaim string
	// DefaultTTL is the default token lifetime used by Mint and /dev/token.
	// Defaults to 1 hour when zero.
	DefaultTTL time.Duration
	// DevStubUser is the sub issued by /authorize. Defaults to "bigbes" when zero.
	DevStubUser string
}

// Stub is an in-memory OIDC stub. Construct with New.
type Stub struct {
	key           *rsa.PrivateKey
	kid           string
	issuer        string
	audience      string
	usernameClaim string
	defaultTTL    time.Duration
	devStubUser   string
	codes         *codeStore
}

// New creates a ready-to-mount Stub. opts.Issuer must be non-empty.
// A fresh 2048-bit RSA signing key is generated on every call (AS3).
func New(opts Options) (*Stub, error) {
	if opts.Issuer == "" {
		return nil, culpa.WithCode(
			culpa.New("oidcstub: Options.Issuer is required"),
			"OIDCSTUB_MISSING_ISSUER",
		)
	}

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, culpa.WithCode(
			culpa.Wrap(err, "oidcstub: generate RSA key"),
			"OIDCSTUB_KEY_GEN",
		)
	}

	ttl := opts.DefaultTTL
	if ttl == 0 {
		ttl = time.Hour
	}

	devUser := opts.DevStubUser
	if devUser == "" {
		devUser = defaultDevStubUser
	}

	return &Stub{
		key:           key,
		kid:           defaultKID,
		issuer:        opts.Issuer,
		audience:      opts.Audience,
		usernameClaim: opts.UsernameClaim,
		defaultTTL:    ttl,
		devStubUser:   devUser,
		codes:         newCodeStore(nil),
	}, nil
}

// SetIssuer updates the issuer URL. Call this once after the httptest.Server
// has started to swap the placeholder URL for the real one. Calling SetIssuer
// after the first token is issued will invalidate already-issued tokens.
func (s *Stub) SetIssuer(issuer string) {
	s.issuer = issuer
}

// Issuer returns the current issuer URL.
func (s *Stub) Issuer() string {
	return s.issuer
}

// Handler returns an http.Handler that serves:
//
//   - /.well-known/openid-configuration — OIDC discovery document
//   - /jwks                             — JSON Web Key Set
//   - /dev/token                        — convenience token endpoint
//   - /authorize                        — auth-code+PKCE authorization endpoint (RFC 6749 §4.1)
//   - /token                            — token endpoint (RFC 6749 §4.1.3, RFC 7636)
func (s *Stub) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/.well-known/openid-configuration", s.handleDiscovery)
	mux.HandleFunc("/jwks", s.handleJWKS)
	mux.HandleFunc("/dev/token", s.handleDevToken)
	mux.HandleFunc("/authorize", s.handleAuthorize)
	mux.HandleFunc("/token", s.handleToken)
	return mux
}

func (s *Stub) handleDiscovery(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	body := map[string]any{
		"issuer":                                s.issuer,
		"jwks_uri":                              s.issuer + "/jwks",
		"authorization_endpoint":                s.issuer + "/authorize",
		"token_endpoint":                        s.issuer + "/token",
		"id_token_signing_alg_values_supported": []string{"RS256"},
		"response_types_supported":              []string{"id_token", "code"},
		"subject_types_supported":               []string{"public"},
		"grant_types_supported":                 []string{"authorization_code"},
		"code_challenge_methods_supported":      []string{"S256"},
	}
	_ = json.NewEncoder(w).Encode(body)
}

func (s *Stub) handleJWKS(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	jwk := map[string]any{
		"kty": "RSA",
		"alg": "RS256",
		"use": "sig",
		"kid": s.kid,
		"n":   base64URLUint(s.key.N),
		"e":   base64URLUintFromInt(s.key.E),
	}
	_ = json.NewEncoder(w).Encode(map[string]any{"keys": []any{jwk}})
}

func (s *Stub) handleDevToken(w http.ResponseWriter, r *http.Request) {
	sub := r.URL.Query().Get("sub")
	if sub == "" {
		http.Error(w, `{"error":"sub is required"}`, http.StatusBadRequest)
		return
	}

	ttl := s.defaultTTL
	if expStr := r.URL.Query().Get("exp"); expStr != "" {
		d, err := time.ParseDuration(expStr)
		if err != nil {
			http.Error(w, `{"error":"invalid exp: `+expStr+`"}`, http.StatusBadRequest)
			return
		}
		ttl = d
	}

	tok, expiresAt, err := s.Mint(sub, ttl, nil)
	if err != nil {
		http.Error(w, `{"error":"mint failed"}`, http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(map[string]any{
		"token":      tok,
		"expires_at": expiresAt.Unix(),
	})
}

// SignToken produces a signed RS256 JWT. The supplied claims are merged with
// defaults {iss, iat, exp=iat+1h}; caller-supplied claims win on conflict.
func (s *Stub) SignToken(claims map[string]any) (string, error) {
	now := time.Now()
	merged := map[string]any{
		"iss": s.issuer,
		"iat": now.Unix(),
		"exp": now.Add(time.Hour).Unix(),
	}
	for k, v := range claims {
		merged[k] = v
	}
	return s.signRaw(merged)
}

// Mint is a convenience method that builds a full token for subject sub.
// It sets: iss, aud (if s.audience is set), iat, exp=iat+ttl,
// the configured UsernameClaim (if non-empty) to sub, and sub=sub.
// Extra claims are merged last so callers can override any default.
// Returns the signed token and its expiration time.
func (s *Stub) Mint(sub string, ttl time.Duration, extra map[string]any) (string, time.Time, error) {
	now := time.Now()
	expiresAt := now.Add(ttl)

	claims := map[string]any{
		"iss": s.issuer,
		"sub": sub,
		"iat": now.Unix(),
		"exp": expiresAt.Unix(),
	}
	if s.audience != "" {
		claims["aud"] = s.audience
	}
	if s.usernameClaim != "" {
		claims[s.usernameClaim] = sub
	}
	// extra overrides everything above.
	for k, v := range extra {
		claims[k] = v
	}

	tok, err := s.signRaw(claims)
	if err != nil {
		return "", time.Time{}, err
	}
	return tok, expiresAt, nil
}

// signRaw signs the supplied claims map exactly as-is (no merging).
func (s *Stub) signRaw(claims map[string]any) (string, error) {
	header := map[string]any{
		"alg": "RS256",
		"typ": "JWT",
		"kid": s.kid,
	}

	headerBytes, err := json.Marshal(header)
	if err != nil {
		return "", culpa.WithCode(
			culpa.Wrap(err, "oidcstub: marshal JWT header"),
			"OIDCSTUB_SIGN",
		)
	}
	payloadBytes, err := json.Marshal(claims)
	if err != nil {
		return "", culpa.WithCode(
			culpa.Wrap(err, "oidcstub: marshal JWT payload"),
			"OIDCSTUB_SIGN",
		)
	}

	signingInput := base64.RawURLEncoding.EncodeToString(headerBytes) +
		"." + base64.RawURLEncoding.EncodeToString(payloadBytes)

	digest := sha256.Sum256([]byte(signingInput))
	sig, err := rsa.SignPKCS1v15(rand.Reader, s.key, crypto.SHA256, digest[:])
	if err != nil {
		return "", culpa.WithCode(
			culpa.Wrap(err, "oidcstub: rsa.SignPKCS1v15"),
			"OIDCSTUB_SIGN",
		)
	}
	return signingInput + "." + base64.RawURLEncoding.EncodeToString(sig), nil
}

// base64URLUint encodes a big.Int as base64url with no padding, per JWA
// (RFC 7518) §6.3.1.1.
func base64URLUint(n *big.Int) string {
	return base64.RawURLEncoding.EncodeToString(n.Bytes())
}

// base64URLUintFromInt encodes the public exponent (typically 65537) as the
// minimal big-endian byte representation, per JWA §6.3.1.2.
func base64URLUintFromInt(e int) string {
	if e <= 0 {
		panic(fmt.Sprintf("oidcstub: invalid RSA exponent: %d", e))
	}
	var buf [8]byte
	binary.BigEndian.PutUint64(buf[:], uint64(e))
	// Strip leading zero bytes so the encoding is minimal.
	i := 0
	for i < len(buf)-1 && buf[i] == 0 {
		i++
	}
	return base64.RawURLEncoding.EncodeToString(buf[i:])
}

// oidcError writes an RFC 6749-compliant JSON error response.
func oidcError(w http.ResponseWriter, status int, code, description string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(map[string]string{
		"error":             code,
		"error_description": description,
	})
}

// handleAuthorize implements the authorization endpoint (RFC 6749 §4.1.1).
// It supports response_type=code with PKCE code_challenge_method=S256 only.
func (s *Stub) handleAuthorize(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		oidcError(w, http.StatusMethodNotAllowed, "invalid_request", "method must be GET")
		return
	}

	q := r.URL.Query()

	responseType := q.Get("response_type")
	if responseType != "code" {
		oidcError(w, http.StatusBadRequest, "invalid_request", "response_type must be 'code'")
		return
	}

	clientID := q.Get("client_id")
	if clientID == "" {
		oidcError(w, http.StatusBadRequest, "invalid_request", "client_id is required")
		return
	}

	redirectURI := q.Get("redirect_uri")
	if redirectURI == "" {
		oidcError(w, http.StatusBadRequest, "invalid_request", "redirect_uri is required")
		return
	}

	state := q.Get("state")
	// state is RECOMMENDED per RFC 6749 §4.1.1; we accept it as optional but echo it.

	codeChallenge := q.Get("code_challenge")
	if codeChallenge == "" {
		oidcError(w, http.StatusBadRequest, "invalid_request", "code_challenge is required")
		return
	}

	codeChallengeMethod := q.Get("code_challenge_method")
	if codeChallengeMethod != "S256" {
		oidcError(w, http.StatusBadRequest, "invalid_request", "code_challenge_method must be 'S256'")
		return
	}

	code := s.codes.Issue(s.devStubUser, codeChallenge, redirectURI, 5*time.Minute)

	// Per RFC 6749 §4.1.2, the redirect Location's query parameters must be
	// properly percent-encoded. The current SPA only emits URL-safe chars,
	// but PC1 requires literal RFC compliance for real-OP parity.
	out := url.Values{"code": {code}}
	if state != "" {
		out.Set("state", state)
	}
	http.Redirect(w, r, redirectURI+"?"+out.Encode(), http.StatusFound)
}

// handleToken implements the token endpoint (RFC 6749 §4.1.3, RFC 7636 §4.6).
// It supports grant_type=authorization_code with PKCE S256 verification.
func (s *Stub) handleToken(w http.ResponseWriter, r *http.Request) {
	// CORS preflight.
	if r.Method == http.MethodOptions {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "POST")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type")
		w.WriteHeader(http.StatusNoContent)
		return
	}

	if r.Method != http.MethodPost {
		oidcError(w, http.StatusMethodNotAllowed, "invalid_request", "method must be POST")
		return
	}

	if err := r.ParseForm(); err != nil {
		oidcError(w, http.StatusBadRequest, "invalid_request", "could not parse form body")
		return
	}

	grantType := r.FormValue("grant_type")
	if grantType != "authorization_code" {
		oidcError(w, http.StatusBadRequest, "unsupported_grant_type", "grant_type must be 'authorization_code'")
		return
	}

	code := r.FormValue("code")
	codeVerifier := r.FormValue("code_verifier")
	redirectURI := r.FormValue("redirect_uri")

	entry, ok := s.codes.Consume(code)
	if !ok {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		oidcError(w, http.StatusBadRequest, "invalid_grant", "unknown or expired code")
		return
	}

	// Verify redirect_uri matches (RFC 6749 §4.1.3).
	if redirectURI != entry.RedirectURI {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		oidcError(w, http.StatusBadRequest, "invalid_grant", "redirect_uri mismatch")
		return
	}

	// Verify PKCE S256: base64url(SHA-256(verifier)) == stored challenge (RFC 7636 §4.6).
	h := sha256.Sum256([]byte(codeVerifier))
	computed := base64.RawURLEncoding.EncodeToString(h[:])
	if computed != entry.CodeChallenge {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		oidcError(w, http.StatusBadRequest, "invalid_grant", "code_verifier does not match code_challenge")
		return
	}

	// Mint JWT.
	tok, _, err := s.Mint(entry.Sub, s.defaultTTL, nil)
	if err != nil {
		http.Error(w, `{"error":"server_error"}`, http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	_ = json.NewEncoder(w).Encode(map[string]any{
		"access_token": tok,
		"id_token":     tok,
		"token_type":   "Bearer",
		"expires_in":   int(s.defaultTTL.Seconds()),
	})
}