package auth

import (
	"context"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/x509"
	"encoding/hex"
	"encoding/pem"
	"errors"
	"fmt"
	"strings"
	"time"

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

// Issuer is "scheduler" — a stable, vendor-namespaced string we emit in our
// self-signed JWTs. The OIDCVerifier uses this to skip OIDC verification on
// our own tokens (saves a JWKS lookup) and the JWTVerifier asserts on it.
const SelfIssuer = "scheduler"

// scheduler-namespaced custom claim keys (ADR-0006 §5.2). Quoted as full URIs
// would be more correct per RFC 7519 but verbose; the spec allows any string
// and the namespace prefix prevents collision with stdlib claim names.
const (
	claimTenantID = "scheduler.v1.tenant_id"
	claimScopes   = "scheduler.v1.scopes"
)

// MaxAPITokenTTL caps how far in the future a self-signed JWT can be valid.
// Operators MUST NOT issue longer tokens — the longer the TTL, the bigger
// the blast radius of a leak.
const MaxAPITokenTTL = 365 * 24 * time.Hour

// DefaultAPITokenTTL is what IssueAPIToken uses if the caller passes 0.
const DefaultAPITokenTTL = 90 * 24 * time.Hour

// Issuer is the signing side. Holds the ES256 private key + a Store reference
// for revocation lifecycle. Built from env via NewIssuerFromEnv on the
// signing node only.
type Issuer struct {
	priv  *ecdsa.PrivateKey
	store Store
	now   func() time.Time // overridable for tests
}

// NewIssuer constructs an Issuer with explicit dependencies. Useful for
// tests; production code should prefer NewIssuerFromEnv.
func NewIssuer(priv *ecdsa.PrivateKey, store Store) *Issuer {
	return &Issuer{priv: priv, store: store, now: time.Now}
}

// NewIssuerFromEnv reads SCHEDULER_JWT_SIGNING_KEY (PEM ES256 private) and
// returns an Issuer. Returns an error if the env var is missing or the key
// is not ES256 — fail-closed configuration per ADR-0006.
func NewIssuerFromEnv(getenv func(string) string, store Store) (*Issuer, error) {
	pemBytes := []byte(getenv("SCHEDULER_JWT_SIGNING_KEY"))
	if len(pemBytes) == 0 {
		return nil, errors.New("auth: SCHEDULER_JWT_SIGNING_KEY not set (deny-by-default)")
	}
	priv, err := parseECPrivateKey(pemBytes)
	if err != nil {
		return nil, fmt.Errorf("auth: parse signing key: %w", err)
	}
	return NewIssuer(priv, store), nil
}

// IssueAPIToken signs a fresh ES256 JWT for the given service account. The
// returned token MUST be handed directly to the operator (typically via the
// CLI stdout) and never written to a log or DB beyond what the test fixture
// needs.
//
// Returns the encoded JWT and its jti so the caller can reference it for
// revocation later.
func (i *Issuer) IssueAPIToken(serviceAccountID, tenantID string, scopes []string, ttl time.Duration) (token string, jti string, err error) {
	if i == nil || i.priv == nil {
		return "", "", errors.New("auth: issuer not configured")
	}
	if serviceAccountID == "" || tenantID == "" {
		return "", "", errors.New("auth: service_account_id and tenant_id required")
	}
	if ttl <= 0 {
		ttl = DefaultAPITokenTTL
	}
	if ttl > MaxAPITokenTTL {
		return "", "", fmt.Errorf("auth: ttl %s exceeds max %s", ttl, MaxAPITokenTTL)
	}
	now := i.now().UTC()
	jtiBytes := make([]byte, 16)
	if _, err := rand.Read(jtiBytes); err != nil {
		return "", "", fmt.Errorf("auth: rand jti: %w", err)
	}
	jti = hex.EncodeToString(jtiBytes)

	// Use MapClaims rather than RegisteredClaims so we can attach the
	// scheduler-namespaced custom keys without a wrapper struct.
	claims := jwt.MapClaims{
		"iss":         SelfIssuer,
		"sub":         serviceAccountID,
		"iat":         now.Unix(),
		"exp":         now.Add(ttl).Unix(),
		"jti":         jti,
		claimTenantID: tenantID,
	}
	if len(scopes) > 0 {
		claims[claimScopes] = scopes
	}

	tok := jwt.NewWithClaims(jwt.SigningMethodES256, claims)
	signed, err := tok.SignedString(i.priv)
	if err != nil {
		return "", "", fmt.Errorf("auth: sign token: %w", err)
	}
	return signed, jti, nil
}

// RevokeAPIToken adds the jti to the denylist with the given expires_at.
// Idempotent.
func (i *Issuer) RevokeAPIToken(ctx context.Context, jti string, expiresAt time.Time) error {
	if i == nil || i.store == nil {
		return errors.New("auth: issuer store not configured")
	}
	return i.store.Revoke(ctx, Revocation{JTI: jti, ExpiresAt: expiresAt})
}

// ---------------------------------------------------------------------------
// JWTVerifier — verifies self-signed scheduler tokens on the receiving side.
// ---------------------------------------------------------------------------

// JWTVerifier verifies self-signed scheduler tokens against the configured
// public key, asserts iss == SelfIssuer, and checks the denylist.
//
// Implements Authenticator.
type JWTVerifier struct {
	pub   *ecdsa.PublicKey
	store Store
	now   func() time.Time
}

// NewJWTVerifier constructs a verifier from an explicit public key. Tests
// pass the pair generated by the same fixture used in jwt_test.go.
func NewJWTVerifier(pub *ecdsa.PublicKey, store Store) *JWTVerifier {
	return &JWTVerifier{pub: pub, store: store, now: time.Now}
}

// NewJWTVerifierFromEnv reads SCHEDULER_JWT_PUBLIC_KEY (PEM ES256 public).
// Like the Issuer constructor, it fails closed on missing/malformed input.
func NewJWTVerifierFromEnv(getenv func(string) string, store Store) (*JWTVerifier, error) {
	pemBytes := []byte(getenv("SCHEDULER_JWT_PUBLIC_KEY"))
	if len(pemBytes) == 0 {
		return nil, errors.New("auth: SCHEDULER_JWT_PUBLIC_KEY not set (deny-by-default)")
	}
	pub, err := parseECPublicKey(pemBytes)
	if err != nil {
		return nil, fmt.Errorf("auth: parse public key: %w", err)
	}
	return NewJWTVerifier(pub, store), nil
}

// Authenticate verifies the bearer token. Returns errAuthenticatorMismatch
// if the token's iss is not SelfIssuer (so ChainAuthenticator can fall
// through to the OIDC verifier).
func (v *JWTVerifier) Authenticate(ctx context.Context, bearer string) (*Actor, error) {
	if v == nil || v.pub == nil {
		return nil, &AuthError{Reason: ReasonAuthnDisabled}
	}
	if bearer == "" {
		return nil, &AuthError{Reason: ReasonMissingCreds}
	}

	// Parse without verification first to check `iss`. If it's not us, fall
	// through. We use jwt.ParseUnverified for this peek; the actual
	// signature check happens below.
	parser := jwt.NewParser(jwt.WithValidMethods([]string{jwt.SigningMethodES256.Alg()}))

	// Inspect alg up-front. RFC 8725: never trust client-supplied alg; we
	// pin to ES256 via WithValidMethods, but a client trying HS256 will
	// get rejected by the parser with a method error — we map that to
	// INSECURE_ALG before falling through.
	_, _, err := parser.ParseUnverified(bearer, jwt.MapClaims{})
	if err != nil {
		// Could be malformed; let the next verifier try.
		return nil, errAuthenticatorMismatch
	}

	tok, err := parser.ParseWithClaims(bearer, jwt.MapClaims{}, func(t *jwt.Token) (any, error) {
		// Belt-and-braces alg check on top of WithValidMethods.
		if t.Method.Alg() != jwt.SigningMethodES256.Alg() {
			return nil, &AuthError{Reason: ReasonInsecureAlg}
		}
		return v.pub, nil
	})

	if err != nil {
		// Did the verify callback bail with our own AuthError?
		var ae *AuthError
		if errors.As(err, &ae) {
			return nil, ae
		}
		// Was iss not SelfIssuer? We need to peek the claims to decide
		// fall-through vs hard-fail.
		if iss := peekIssuer(bearer); iss != "" && iss != SelfIssuer {
			return nil, errAuthenticatorMismatch
		}
		// Distinguish expired from "everything else" so the audit hook
		// records the precise reason. The client-visible message stays
		// uniform — that is enforced by the middleware, NOT here.
		if errors.Is(err, jwt.ErrTokenExpired) {
			return nil, &AuthError{Reason: ReasonTokenExpired, Err: err}
		}
		return nil, &AuthError{Reason: ReasonInvalidToken, Err: err}
	}

	claims, ok := tok.Claims.(jwt.MapClaims)
	if !ok || !tok.Valid {
		return nil, &AuthError{Reason: ReasonInvalidToken}
	}

	iss, _ := claims["iss"].(string)
	if iss != SelfIssuer {
		return nil, errAuthenticatorMismatch
	}

	// Check exp explicitly — golang-jwt/v5 already validates it via the
	// parser by default, but doing it here keeps the reason mapping
	// deterministic.
	if expF, ok := claims["exp"].(float64); ok {
		if v.now().Unix() >= int64(expF) {
			return nil, &AuthError{Reason: ReasonTokenExpired}
		}
	}

	tenantID, _ := claims[claimTenantID].(string)
	if tenantID == "" {
		return nil, &AuthError{Reason: ReasonMissingTenant}
	}
	sub, _ := claims["sub"].(string)
	if sub == "" {
		return nil, &AuthError{Reason: ReasonInvalidToken}
	}
	jti, _ := claims["jti"].(string)
	if jti != "" && v.store != nil {
		revoked, err := v.store.IsRevoked(ctx, jti)
		if err != nil {
			return nil, &AuthError{Reason: ReasonInvalidToken, Err: err}
		}
		if revoked {
			return nil, &AuthError{Reason: ReasonTokenRevoked}
		}
	}

	scopes := parseScopesClaim(claims[claimScopes])

	return &Actor{
		Type:     ActorService,
		ID:       sub,
		TenantID: tenantID,
		Scopes:   scopes,
	}, nil
}

// peekIssuer parses the JWT WITHOUT verifying signature, just to look at the
// `iss` claim. Used only to decide whether to fall through to OIDC.
func peekIssuer(bearer string) string {
	parts := strings.Split(bearer, ".")
	if len(parts) != 3 {
		return ""
	}
	parser := jwt.NewParser()
	tok, _, err := parser.ParseUnverified(bearer, jwt.MapClaims{})
	if err != nil {
		return ""
	}
	if c, ok := tok.Claims.(jwt.MapClaims); ok {
		if iss, ok := c["iss"].(string); ok {
			return iss
		}
	}
	return ""
}

func parseScopesClaim(v any) []string {
	switch t := v.(type) {
	case []any:
		out := make([]string, 0, len(t))
		for _, e := range t {
			if s, ok := e.(string); ok && s != "" {
				out = append(out, s)
			}
		}
		return out
	case []string:
		return append([]string(nil), t...)
	default:
		return nil
	}
}

// ---------------------------------------------------------------------------
// PEM helpers.
// ---------------------------------------------------------------------------

func parseECPrivateKey(pemBytes []byte) (*ecdsa.PrivateKey, error) {
	block, _ := pem.Decode(pemBytes)
	if block == nil {
		return nil, errors.New("auth: not PEM")
	}
	// PKCS8 first, fall back to SEC1.
	if k, err := x509.ParsePKCS8PrivateKey(block.Bytes); err == nil {
		ec, ok := k.(*ecdsa.PrivateKey)
		if !ok {
			return nil, errors.New("auth: signing key not ECDSA")
		}
		return ec, nil
	}
	ec, err := x509.ParseECPrivateKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("auth: parse EC key: %w", err)
	}
	return ec, nil
}

func parseECPublicKey(pemBytes []byte) (*ecdsa.PublicKey, error) {
	block, _ := pem.Decode(pemBytes)
	if block == nil {
		return nil, errors.New("auth: not PEM")
	}
	k, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("auth: parse PKIX: %w", err)
	}
	ec, ok := k.(*ecdsa.PublicKey)
	if !ok {
		return nil, errors.New("auth: public key not ECDSA")
	}
	return ec, nil
}