package wristband

import (
	"context"
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/pem"
	"errors"
	"fmt"
	"math/big"
	"time"

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

// The id_token signing algorithms (D-12).
const (
	// AlgRS256 is RSASSA-PKCS1-v1_5 with SHA-256 over an RSA key, the
	// default and the algorithm of imported PHP keys.
	AlgRS256 = "RS256"
	// AlgES256 is ECDSA over P-256 with SHA-256.
	AlgES256 = "ES256"
)

// SigningKey is the active signing key. Signer must be the
// *rsa.PrivateKey or *ecdsa.PrivateKey matching Alg, as returned by
// ParsePrivateKeyPEM. Key material reaches the engine only here and is
// never logged (T-19.4-09).
type SigningKey struct {
	Kid    string
	Alg    string
	Signer crypto.Signer
}

// PublicKey is a published verification key.
type PublicKey struct {
	Kid string
	Alg string
	Key crypto.PublicKey
}

// KeyStore supplies the signing-key ring (D-10..D-13).
type KeyStore interface {
	// ActiveSigningKey returns the active key, generating the first key
	// when the ring is empty.
	ActiveSigningKey(ctx context.Context) (SigningKey, error)
	// PublishedKeys returns the active key plus retired keys still inside
	// their grace window at now, for JWKS and discovery.
	PublishedKeys(ctx context.Context, now time.Time) ([]PublicKey, error)
}

// GenerateSigningKey creates a new key for alg: RSA 2048 for RS256, P-256
// for ES256; any other algorithm is refused. It
// returns a kid of 8 random bytes in lowercase hex (16 characters, like the
// PHP key ring), the private key as PKCS#8 PEM and the public key as PKIX
// PEM.
func GenerateSigningKey(alg string) (kid string, privatePEM, publicPEM []byte, err error) {
	var priv crypto.Signer
	switch alg {
	case AlgRS256:
		priv, err = rsa.GenerateKey(rand.Reader, 2048)
	case AlgES256:
		priv, err = ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	default:
		return "", nil, nil, fmt.Errorf("wristband: unsupported signing algorithm %q", alg)
	}
	if err != nil {
		return "", nil, nil, fmt.Errorf("wristband: generate signing key: %w", err)
	}
	privDER, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return "", nil, nil, fmt.Errorf("wristband: encode private key: %w", err)
	}
	pubDER, err := x509.MarshalPKIXPublicKey(priv.Public())
	if err != nil {
		return "", nil, nil, fmt.Errorf("wristband: encode public key: %w", err)
	}
	raw := make([]byte, 8)
	if _, err := rand.Read(raw); err != nil {
		return "", nil, nil, fmt.Errorf("wristband: generate kid: %w", err)
	}
	privatePEM = pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: privDER})
	publicPEM = pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubDER})
	return hex.EncodeToString(raw), privatePEM, publicPEM, nil
}

// ParsePrivateKeyPEM parses a private key: PKCS#8 first, then PKCS#1 RSA,
// then SEC1 EC. It returns an *rsa.PrivateKey or a P-256
// *ecdsa.PrivateKey; errors never contain key material (T-19.4-09).
func ParsePrivateKeyPEM(data []byte) (crypto.Signer, error) {
	block, _ := pem.Decode(data)
	if block == nil {
		return nil, errors.New("wristband: no PEM block in private key")
	}
	var key any
	var err error
	if key, err = x509.ParsePKCS8PrivateKey(block.Bytes); err != nil {
		if key, err = x509.ParsePKCS1PrivateKey(block.Bytes); err != nil {
			if key, err = x509.ParseECPrivateKey(block.Bytes); err != nil {
				return nil, errors.New("wristband: unreadable private key")
			}
		}
	}
	switch k := key.(type) {
	case *rsa.PrivateKey:
		return k, nil
	case *ecdsa.PrivateKey:
		if k.Curve == elliptic.P256() {
			return k, nil
		}
	}
	return nil, errors.New("wristband: unsupported private key type")
}

// ParsePublicKeyPEM parses a PKIX public key (RSA or P-256 EC).
func ParsePublicKeyPEM(data []byte) (crypto.PublicKey, error) {
	block, _ := pem.Decode(data)
	if block == nil {
		return nil, errors.New("wristband: no PEM block in public key")
	}
	key, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, errors.New("wristband: unreadable public key")
	}
	switch k := key.(type) {
	case *rsa.PublicKey:
		return k, nil
	case *ecdsa.PublicKey:
		if k.Curve == elliptic.P256() {
			return k, nil
		}
	}
	return nil, errors.New("wristband: unsupported public key type")
}

// jwkDocument is one JWKS entry. encoding/json keeps declaration order and
// omitempty drops the members of the other key type, so RSA keys serialise
// as kty, use, alg, kid, n, e and EC keys as kty, use, alg, kid, crv, x, y.
type jwkDocument struct {
	Kty string `json:"kty"`
	Use string `json:"use"`
	Alg string `json:"alg"`
	Kid string `json:"kid"`
	Crv string `json:"crv,omitempty"`
	N   string `json:"n,omitempty"`
	E   string `json:"e,omitempty"`
	X   string `json:"x,omitempty"`
	Y   string `json:"y,omitempty"`
}

// jwkFor encodes a published key with the standard library (RFC 7518
// §6.3.1: n and e are minimal big-endian bytes; §6.2.1: x and y are the
// fixed-length halves of the uncompressed P-256 point from
// ecdsa.PublicKey.Bytes, so no padding bug is possible). Every JWK carries
// use sig
// and its alg; an Alg that does not match the key type is refused
// (T-19.4-08).
func jwkFor(k PublicKey) (jwkDocument, error) {
	b64 := base64.RawURLEncoding.EncodeToString
	switch key := k.Key.(type) {
	case *rsa.PublicKey:
		if k.Alg != AlgRS256 {
			return jwkDocument{}, errors.New("wristband: RSA key with a non-RSA algorithm")
		}
		return jwkDocument{
			Kty: "RSA", Use: "sig", Alg: k.Alg, Kid: k.Kid,
			N: b64(key.N.Bytes()), E: b64(big.NewInt(int64(key.E)).Bytes()),
		}, nil
	case *ecdsa.PublicKey:
		if k.Alg != AlgES256 || key.Curve != elliptic.P256() {
			return jwkDocument{}, errors.New("wristband: EC key that is not P-256 with ES256")
		}
		raw, err := key.Bytes()
		if err != nil || len(raw) != 65 {
			return jwkDocument{}, errors.New("wristband: unsupported EC key")
		}
		return jwkDocument{
			Kty: "EC", Use: "sig", Alg: k.Alg, Kid: k.Kid,
			Crv: "P-256", X: b64(raw[1:33]), Y: b64(raw[33:65]),
		}, nil
	}
	return jwkDocument{}, errors.New("wristband: unsupported public key type")
}

// idTokenInput carries the claims of one id_token (PHP TokenIssuer).
type idTokenInput struct {
	Issuer      string
	Subject     string
	Audience    string
	TTL         time.Duration
	Now         time.Time
	Nonce       string
	AccessToken string
	Identity    map[string]any
	Allowed     []string
}

// protocolClaims are the claims the engine owns; an identity claim with one
// of these names is never copied (T-19.4-11).
var protocolClaims = map[string]bool{
	"iss": true, "sub": true, "aud": true, "exp": true, "iat": true,
	"nonce": true, "at_hash": true, "auth_time": true, "azp": true, "nbf": true, "jti": true,
}

// signIDToken signs an id_token with golang-jwt map claims, so aud stays a
// JSON string as in the PHP provider (jwt.RegisteredClaims would marshal it
// as an array under a package-global default that bouncer shares and that
// wristband never changes). Allowlisted identity claims are
// written first, then the protocol claims, then sub once more, so a user
// attribute can never overwrite a protocol claim (T-19.4-11).
func signIDToken(key SigningKey, in idTokenInput) (string, error) {
	method, err := signingMethodFor(key)
	if err != nil {
		return "", err
	}
	claims := jwt.MapClaims{}
	for _, name := range in.Allowed {
		if protocolClaims[name] {
			continue
		}
		if v, ok := in.Identity[name]; ok {
			claims[name] = v
		}
	}
	claims["iss"] = in.Issuer
	claims["sub"] = in.Subject
	claims["aud"] = in.Audience
	claims["exp"] = in.Now.Add(in.TTL).Unix()
	claims["iat"] = in.Now.Unix()
	if in.Nonce != "" {
		claims["nonce"] = in.Nonce
	}
	if in.AccessToken != "" {
		claims["at_hash"] = atHash(in.AccessToken)
	}
	claims["sub"] = in.Subject

	tok := jwt.NewWithClaims(method, claims)
	tok.Header["kid"] = key.Kid
	signed, err := tok.SignedString(key.Signer)
	if err != nil {
		return "", errors.New("wristband: id_token signing failed")
	}
	return signed, nil
}

// signingMethodFor derives the JWS method from the key type and refuses a
// key whose Alg does not match it (T-19.4-08).
func signingMethodFor(key SigningKey) (jwt.SigningMethod, error) {
	switch key.Alg {
	case AlgRS256:
		if _, ok := key.Signer.(*rsa.PrivateKey); ok {
			return jwt.SigningMethodRS256, nil
		}
	case AlgES256:
		if k, ok := key.Signer.(*ecdsa.PrivateKey); ok && k.Curve == elliptic.P256() {
			return jwt.SigningMethodES256, nil
		}
	}
	return nil, errors.New("wristband: signing key does not match its algorithm")
}

// atHash is base64url of the left half of SHA-256 over the access token
// (OIDC Core §3.1.3.6; RS256 and ES256 both hash with SHA-256).
func atHash(accessToken string) string {
	sum := sha256.Sum256([]byte(accessToken))
	return base64.RawURLEncoding.EncodeToString(sum[:16])
}