// Package oidctest provides a minimal OIDC provider for tests —
// discovery doc, JWKS, and a token endpoint. Each test mints its
// own claims; the stub signs them with an ECDSA P-256 key and the
// production verifier accepts them because the JWKS is fetched live
// from the stub.
//
// Usage:
//
//	stub := oidctest.New(t)
//	code := stub.MintCode(map[string]any{
//	    "sub":                "abc",
//	    "preferred_username": "alice",
//	    "groups":             []string{"rm-admins"},
//	})
//	// stub.URL() is the issuer URL; pass to oidc.New as Issuer
package oidctest

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	stdhttp "net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

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

// StubIdP is an httptest-backed OIDC provider. Each test creates a
// fresh one via New(t); cleanup is registered on t.
type StubIdP struct {
	t   *testing.T
	srv *httptest.Server

	mu         sync.Mutex
	priv       *ecdsa.PrivateKey
	kid        string
	claims     map[string]map[string]any // code → claims
	endSession string                    // optional, set by SetEndSessionEndpoint
}

// New constructs a stub IdP listening on a random port. Cleanup is
// registered on t.
func New(t *testing.T) *StubIdP {
	t.Helper()
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("oidctest: genkey: %v", err)
	}
	s := &StubIdP{
		t:      t,
		priv:   priv,
		kid:    "stub-key",
		claims: map[string]map[string]any{},
	}
	mux := stdhttp.NewServeMux()
	mux.HandleFunc("/.well-known/openid-configuration", s.discovery)
	mux.HandleFunc("/jwks.json", s.jwks)
	mux.HandleFunc("/token", s.token)
	s.srv = httptest.NewServer(mux)
	t.Cleanup(s.srv.Close)
	return s
}

// URL returns the base URL of the stub — pass as Issuer to
// oidc.New().
func (s *StubIdP) URL() string { return s.srv.URL }

// MintCode produces an authorization code that the stub will exchange
// for an id_token containing the supplied claims.
func (s *StubIdP) MintCode(claims map[string]any) string {
	s.mu.Lock()
	defer s.mu.Unlock()
	code := fmt.Sprintf("code-%d", time.Now().UnixNano())
	s.claims[code] = claims
	return code
}

// SetEndSessionEndpoint configures the stub to advertise an
// end_session_endpoint in its discovery doc. Used by the logout
// test in E1.
func (s *StubIdP) SetEndSessionEndpoint(url string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.endSession = url
}

func (s *StubIdP) discovery(w stdhttp.ResponseWriter, _ *stdhttp.Request) {
	s.mu.Lock()
	endSession := s.endSession
	s.mu.Unlock()

	doc := map[string]any{
		"issuer":                                s.srv.URL,
		"authorization_endpoint":                s.srv.URL + "/authorize",
		"token_endpoint":                        s.srv.URL + "/token",
		"jwks_uri":                              s.srv.URL + "/jwks.json",
		"id_token_signing_alg_values_supported": []string{"ES256"},
		"response_types_supported":              []string{"code"},
		"subject_types_supported":               []string{"public"},
	}
	if endSession != "" {
		doc["end_session_endpoint"] = endSession
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(doc)
}

func (s *StubIdP) jwks(w stdhttp.ResponseWriter, _ *stdhttp.Request) {
	pub := s.priv.Public().(*ecdsa.PublicKey)
	x := base64.RawURLEncoding.EncodeToString(padTo32(pub.X.Bytes()))
	y := base64.RawURLEncoding.EncodeToString(padTo32(pub.Y.Bytes()))
	keys := map[string]any{
		"keys": []map[string]any{{
			"kty": "EC", "crv": "P-256", "alg": "ES256",
			"use": "sig", "kid": s.kid,
			"x": x, "y": y,
		}},
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(keys)
}

func (s *StubIdP) token(w stdhttp.ResponseWriter, r *stdhttp.Request) {
	_ = r.ParseForm()
	code := r.PostForm.Get("code")
	s.mu.Lock()
	claims, ok := s.claims[code]
	if ok {
		delete(s.claims, code)
	}
	s.mu.Unlock()
	if !ok {
		stdhttp.Error(w, "bad code", stdhttp.StatusBadRequest)
		return
	}
	if _, ok := claims["iss"]; !ok {
		claims["iss"] = s.srv.URL
	}
	if _, ok := claims["aud"]; !ok {
		claims["aud"] = "test-client"
	}
	now := time.Now().Unix()
	claims["iat"] = now
	claims["exp"] = now + 600

	jc := jwt.MapClaims{}
	for k, v := range claims {
		jc[k] = v
	}
	tk := jwt.NewWithClaims(jwt.SigningMethodES256, jc)
	tk.Header["kid"] = s.kid
	signed, err := tk.SignedString(s.priv)
	if err != nil {
		stdhttp.Error(w, err.Error(), 500)
		return
	}
	resp := map[string]any{
		"access_token": "stub-access",
		"token_type":   "Bearer",
		"id_token":     signed,
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(resp)
}

// padTo32 left-pads an integer big-endian byte slice to 32 bytes,
// the size required by P-256 JWK x/y components.
func padTo32(b []byte) []byte {
	if len(b) >= 32 {
		return b
	}
	out := make([]byte, 32)
	copy(out[32-len(b):], b)
	return out
}