// 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"
	"time"

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

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
}

// 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
}

// 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
	}

	return &Stub{
		key:           key,
		kid:           defaultKID,
		issuer:        opts.Issuer,
		audience:      opts.Audience,
		usernameClaim: opts.UsernameClaim,
		defaultTTL:    ttl,
	}, 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
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)
	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 + "/auth",
		"token_endpoint":                        s.issuer + "/token",
		"id_token_signing_alg_values_supported": []string{"RS256"},
		"response_types_supported":              []string{"id_token"},
		"subject_types_supported":               []string{"public"},
	}
	_ = 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:])
}