package auth

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"net/http"
	"slices"
	"strings"
	"sync"
	"time"
)

// OIDC validates `Authorization: Bearer <jwt>` against any OIDC issuer: it
// fetches `<issuer>/.well-known/openid-configuration` once, caches the JWKS,
// and verifies signature (RS256/ES256), iss, aud and exp/nbf. Zitadel is the
// reference issuer but nothing here is provider-specific.
//
// Note: the issuer must mint JWT access tokens (for Zitadel, set the
// application's access-token type to JWT); opaque tokens require
// introspection, which this deliberately does not implement.
type OIDC struct {
	issuer   string
	audience string // optional: when set, must appear in the token's aud
	client   *http.Client

	mu          sync.Mutex
	jwksURI     string
	keys        map[string]crypto.PublicKey // kid → verification key
	lastRefresh time.Time

	now func() time.Time // injectable clock for tests
}

var _ Authenticator = (*OIDC)(nil)

// NewOIDC builds an OIDC authenticator for the given issuer URL. audience
// is the expected `aud` claim (typically the OAuth client id); empty skips
// the audience check.
func NewOIDC(issuer, audience string) *OIDC {
	return &OIDC{
		issuer:   strings.TrimSuffix(issuer, "/"),
		audience: audience,
		client:   &http.Client{Timeout: 10 * time.Second},
		now:      time.Now,
	}
}

// Authenticate validates the bearer JWT on r and returns the caller as a
// minimal Principal. For callers that need the full parsed claim set
// (Email/Name/scope list) — e.g. admin handlers rendering who is logged in —
// see VerifyClaims.
func (o *OIDC) Authenticate(r *http.Request) (*Principal, error) {
	h := r.Header.Get("Authorization")
	token, ok := strings.CutPrefix(h, "Bearer ")
	if !ok || token == "" {
		return nil, ErrNoCredentials
	}
	c, err := o.VerifyClaims(token)
	if err != nil {
		return nil, err
	}
	return &Principal{Subject: c.Subject, Scopes: c.Scopes, Roles: c.Roles, Kind: KindOIDC}, nil
}

// Claims is the parsed, verified claim set returned by VerifyClaims. Only
// the subset services commonly need is exposed; the raw claim JSON is
// available via the original token, which callers can re-decode if they
// need something custom.
type Claims struct {
	Issuer    string   `json:"iss"`
	Subject   string   `json:"sub"`
	Email     string   `json:"email,omitempty"`
	Name      string   `json:"name,omitempty"`
	Audience  Audience `json:"aud"`
	Expires   int64    `json:"exp"`
	NotBefore int64    `json:"nbf"`
	Scopes    []string `json:"scope,omitempty"` // parsed from "scope" space-separated claim
	Roles     []string `json:"roles,omitempty"` // parsed from "roles" or "groups" claim
}

// Audience unmarshals the `aud` claim, which RFC 7519 allows to be either a
// string or an array of strings (some providers emit one, some the other;
// Zitadel emits an array).
type Audience []string

func (a *Audience) UnmarshalJSON(b []byte) error {
	var s string
	if err := json.Unmarshal(b, &s); err == nil {
		*a = Audience{s}
		return nil
	}
	var ss []string
	if err := json.Unmarshal(b, &ss); err != nil {
		return err
	}
	*a = Audience(ss)
	return nil
}

// contains reports whether want appears in the audience list.
func (a Audience) contains(want string) bool {
	return slices.Contains(a, want)
}

// clockSkew is the tolerance applied when checking expiry/notBefore so a
// small clock drift between issuer and verifier doesn't reject a fresh
// token.
const clockSkew = 30 * time.Second

// VerifyClaims verifies a raw JWT (no "Bearer " prefix) end-to-end and
// returns the parsed claim set. Callers that aren't using the standard
// Authorization header (e.g. an admin UI reading a session cookie) can
// pass the cookie value directly. The token is verified for issuer,
// audience (if configured), expiry/notBefore, and signature via the JWKS
// cached against the configured issuer.
func (o *OIDC) VerifyClaims(token string) (*Claims, error) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return nil, errors.New("oidc: malformed jwt")
	}
	headerJSON, err := base64.RawURLEncoding.DecodeString(parts[0])
	if err != nil {
		return nil, fmt.Errorf("oidc: bad jwt header: %w", err)
	}
	var header struct {
		Alg string `json:"alg"`
		Kid string `json:"kid"`
	}
	if err := json.Unmarshal(headerJSON, &header); err != nil {
		return nil, fmt.Errorf("oidc: bad jwt header: %w", err)
	}
	payloadJSON, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return nil, fmt.Errorf("oidc: bad jwt payload: %w", err)
	}
	sig, err := base64.RawURLEncoding.DecodeString(parts[2])
	if err != nil {
		return nil, fmt.Errorf("oidc: bad jwt signature: %w", err)
	}

	key, err := o.key(header.Kid)
	if err != nil {
		return nil, err
	}
	digest := sha256.Sum256([]byte(parts[0] + "." + parts[1]))
	switch header.Alg {
	case "RS256":
		pub, ok := key.(*rsa.PublicKey)
		if !ok {
			return nil, errors.New("oidc: key/alg mismatch")
		}
		if err := rsa.VerifyPKCS1v15(pub, crypto.SHA256, digest[:], sig); err != nil {
			return nil, errors.New("oidc: invalid signature")
		}
	case "ES256":
		pub, ok := key.(*ecdsa.PublicKey)
		if !ok {
			return nil, errors.New("oidc: key/alg mismatch")
		}
		if len(sig) != 64 {
			return nil, errors.New("oidc: invalid signature")
		}
		rr := new(big.Int).SetBytes(sig[:32])
		ss := new(big.Int).SetBytes(sig[32:])
		if !ecdsa.Verify(pub, digest[:], rr, ss) {
			return nil, errors.New("oidc: invalid signature")
		}
	default:
		return nil, fmt.Errorf("oidc: unsupported alg %q", header.Alg)
	}

	// Decode into a wide struct so Email/Name are picked up regardless of
	// claim shape (some providers omit them, some populate them).
	var raw struct {
		Issuer    string   `json:"iss"`
		Subject   string   `json:"sub"`
		Email     string   `json:"email"`
		Name      string   `json:"name"`
		Audience  Audience `json:"aud"`
		Expires   int64    `json:"exp"`
		NotBefore int64    `json:"nbf"`
		Scope     string   `json:"scope"`
		Roles     []string `json:"roles"`
		Groups    []string `json:"groups"`
	}
	if err := json.Unmarshal(payloadJSON, &raw); err != nil {
		return nil, fmt.Errorf("oidc: bad claims: %w", err)
	}
	now := o.now()
	if strings.TrimSuffix(raw.Issuer, "/") != o.issuer {
		return nil, errors.New("oidc: wrong issuer")
	}
	if raw.Expires != 0 && now.After(time.Unix(raw.Expires, 0).Add(clockSkew)) {
		return nil, errors.New("oidc: token expired")
	}
	if raw.NotBefore != 0 && now.Add(clockSkew).Before(time.Unix(raw.NotBefore, 0)) {
		return nil, errors.New("oidc: token not yet valid")
	}
	if o.audience != "" && !raw.Audience.contains(o.audience) {
		return nil, errors.New("oidc: wrong audience")
	}

	var scopes []string
	if raw.Scope != "" {
		scopes = strings.Fields(raw.Scope)
	}
	roles := raw.Roles
	if len(roles) == 0 && len(raw.Groups) > 0 {
		roles = raw.Groups
	}

	return &Claims{
		Issuer:    raw.Issuer,
		Subject:   raw.Subject,
		Email:     raw.Email,
		Name:      raw.Name,
		Audience:  raw.Audience,
		Expires:   raw.Expires,
		NotBefore: raw.NotBefore,
		Scopes:    scopes,
		Roles:     roles,
	}, nil
}

// key returns the verification key for kid, fetching/refreshing the JWKS
// when the kid is unknown (key rotation), rate-limited to once a minute.
func (o *OIDC) key(kid string) (crypto.PublicKey, error) {
	o.mu.Lock()
	defer o.mu.Unlock()
	if k, ok := o.keys[kid]; ok {
		return k, nil
	}
	if time.Since(o.lastRefresh) < time.Minute {
		return nil, fmt.Errorf("oidc: unknown key id %q", kid)
	}
	if err := o.refreshLocked(); err != nil {
		return nil, err
	}
	if k, ok := o.keys[kid]; ok {
		return k, nil
	}
	return nil, fmt.Errorf("oidc: unknown key id %q", kid)
}

// refreshLocked (re)fetches the discovery document (first time) and the
// JWKS. Callers must hold o.mu.
func (o *OIDC) refreshLocked() error {
	o.lastRefresh = o.now()
	if o.jwksURI == "" {
		var disc struct {
			JWKSURI string `json:"jwks_uri"`
		}
		if err := o.getJSON(o.issuer+"/.well-known/openid-configuration", &disc); err != nil {
			return fmt.Errorf("oidc: discovery: %w", err)
		}
		if disc.JWKSURI == "" {
			return errors.New("oidc: discovery document has no jwks_uri")
		}
		o.jwksURI = disc.JWKSURI
	}
	var set struct {
		Keys []jwk `json:"keys"`
	}
	if err := o.getJSON(o.jwksURI, &set); err != nil {
		return fmt.Errorf("oidc: jwks: %w", err)
	}
	keys := make(map[string]crypto.PublicKey, len(set.Keys))
	for _, k := range set.Keys {
		pub, err := k.publicKey()
		if err != nil {
			continue // skip unsupported key types (e.g. enc-use keys)
		}
		keys[k.Kid] = pub
	}
	o.keys = keys
	return nil
}

func (o *OIDC) getJSON(url string, v any) error {
	res, err := o.client.Get(url)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("GET %s: %s", url, res.Status)
	}
	return json.NewDecoder(res.Body).Decode(v)
}

// jwk is the subset of RFC 7517 needed to build RSA and P-256 keys.
type jwk struct {
	Kty string `json:"kty"`
	Kid string `json:"kid"`
	Use string `json:"use"`
	N   string `json:"n"`
	E   string `json:"e"`
	Crv string `json:"crv"`
	X   string `json:"x"`
	Y   string `json:"y"`
}

func (k jwk) publicKey() (crypto.PublicKey, error) {
	switch k.Kty {
	case "RSA":
		n, err := base64.RawURLEncoding.DecodeString(k.N)
		if err != nil {
			return nil, err
		}
		eBytes, err := base64.RawURLEncoding.DecodeString(k.E)
		if err != nil {
			return nil, err
		}
		return &rsa.PublicKey{
			N: new(big.Int).SetBytes(n),
			E: int(new(big.Int).SetBytes(eBytes).Int64()),
		}, nil
	case "EC":
		if k.Crv != "P-256" {
			return nil, fmt.Errorf("unsupported curve %q", k.Crv)
		}
		x, err := base64.RawURLEncoding.DecodeString(k.X)
		if err != nil {
			return nil, err
		}
		y, err := base64.RawURLEncoding.DecodeString(k.Y)
		if err != nil {
			return nil, err
		}
		return &ecdsa.PublicKey{
			Curve: elliptic.P256(),
			X:     new(big.Int).SetBytes(x),
			Y:     new(big.Int).SetBytes(y),
		}, nil
	default:
		return nil, fmt.Errorf("unsupported key type %q", k.Kty)
	}
}