package oidc

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"math/big"
	"net/http"
)

// JWK represents a single JSON Web Key for an RSA public key.
// Only fields required for RS256 signature verification are included.
type JWK struct {
	Kty string `json:"kty"` // "RSA"
	Use string `json:"use"` // "sig"
	Alg string `json:"alg"` // "RS256"
	Kid string `json:"kid"` // key identifier
	N   string `json:"n"`   // base64url-encoded modulus (no padding)
	E   string `json:"e"`   // base64url-encoded public exponent (no padding)
}

// keyEntry pairs a kid with the corresponding public key.
type keyEntry struct {
	kid string
	pub *rsa.PublicKey
}

// KeySet holds one or more RSA public keys for JWKS rotation.
// Keys are served in insertion order.
type KeySet struct {
	entries []keyEntry
}

// NewKeySet returns an empty KeySet ready for AddKey calls.
func NewKeySet() *KeySet {
	return &KeySet{}
}

// AddKey appends an RSA public key with the given key ID.
// kid must be unique within the set; duplicates are not checked.
func (ks *KeySet) AddKey(kid string, pub *rsa.PublicKey) {
	ks.entries = append(ks.entries, keyEntry{kid: kid, pub: pub})
}

// jwkFromPublicKey encodes an RSA public key as a JWK using base64url (no padding).
func jwkFromPublicKey(kid string, pub *rsa.PublicKey) JWK {
	enc := base64.RawURLEncoding

	// Modulus — big-endian bytes, no leading zero (math/big ensures minimal encoding).
	nBytes := pub.N.Bytes()

	// Exponent — big-endian minimal encoding.
	exp := big.NewInt(int64(pub.E))
	eBytes := exp.Bytes()

	return JWK{
		Kty: "RSA",
		Use: "sig",
		Alg: "RS256",
		Kid: kid,
		N:   enc.EncodeToString(nBytes),
		E:   enc.EncodeToString(eBytes),
	}
}

// jwksResponse is the top-level JWK Set object.
type jwksResponse struct {
	Keys []JWK `json:"keys"`
}

// jwksHandler implements http.Handler for GET /jwks.
type jwksHandler struct {
	ks *KeySet
}

// NewJWKSHandler returns an http.Handler that serves the JWK Set.
// The key set is serialised on every request so key rotation can be supported
// by mutating the KeySet before the next request (safe for construction-time use;
// for live rotation a RWMutex should wrap AddKey).
func NewJWKSHandler(ks *KeySet) http.Handler {
	return &jwksHandler{ks: ks}
}

func (h *jwksHandler) ServeHTTP(w http.ResponseWriter, _ *http.Request) {
	jwks := jwksResponse{Keys: make([]JWK, 0, len(h.ks.entries))}
	for _, e := range h.ks.entries {
		jwks.Keys = append(jwks.Keys, jwkFromPublicKey(e.kid, e.pub))
	}

	b, err := json.Marshal(jwks)
	if err != nil {
		http.Error(w, "internal error encoding JWKS", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write(b)
}

// LoadPublicKeyFromPEM parses a PEM-encoded public key (PKIX / "PUBLIC KEY" block).
// Returns an error if the PEM data is invalid or does not contain an RSA public key.
func LoadPublicKeyFromPEM(pemData []byte) (*rsa.PublicKey, error) {
	block, _ := pem.Decode(pemData)
	if block == nil {
		return nil, errors.New("jwks: no PEM block found in input")
	}
	if block.Type != "PUBLIC KEY" {
		return nil, errors.New("jwks: expected PEM block type \"PUBLIC KEY\", got \"" + block.Type + "\"")
	}
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, errors.New("jwks: failed to parse PKIX public key: " + err.Error())
	}
	rsaPub, ok := pub.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("jwks: key is not an RSA public key")
	}
	return rsaPub, nil
}