package auth

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"os"
	"strings"

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

// JWTConfig holds RS256 validation settings (local PEM now; JWKS URL later for Okta).
type JWTConfig struct {
	Enabled   bool
	Required  bool
	Issuer    string
	Audience  string
	PublicKey *rsa.PublicKey
}

// LoadPublicKeyFromPEM parses an RSA public key from PEM bytes.
func LoadPublicKeyFromPEM(pemBytes []byte) (*rsa.PublicKey, error) {
	block, _ := pem.Decode(pemBytes)
	if block == nil {
		return nil, errors.New("invalid PEM")
	}
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	rsaPub, ok := pub.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("not an RSA public key")
	}
	return rsaPub, nil
}

// LoadPublicKeyFromFile reads a PEM public key from disk.
func LoadPublicKeyFromFile(path string) (*rsa.PublicKey, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	return LoadPublicKeyFromPEM(b)
}

// Validator verifies Bearer JWTs.
type Validator struct {
	cfg JWTConfig
}

// NewValidator constructs a JWT validator.
func NewValidator(cfg JWTConfig) (*Validator, error) {
	if !cfg.Enabled {
		return &Validator{cfg: cfg}, nil
	}
	if cfg.PublicKey == nil {
		return nil, errors.New("JWT enabled but no public key configured")
	}
	return &Validator{cfg: cfg}, nil
}

// Enabled reports whether JWT validation is active.
func (v *Validator) Enabled() bool {
	return v.cfg.Enabled
}

// Required reports whether requests must include a valid Bearer token.
func (v *Validator) Required() bool {
	return v.cfg.Required
}

// Validate parses and verifies a Bearer token string (without "Bearer " prefix).
func (v *Validator) Validate(tokenString string) (*Principal, error) {
	if !v.cfg.Enabled {
		return nil, nil
	}
	tokenString = strings.TrimSpace(tokenString)
	if tokenString == "" {
		if v.cfg.Required {
			return nil, errors.New("missing bearer token")
		}
		return nil, nil
	}

	keyFunc := func(t *jwt.Token) (any, error) {
		if t.Method.Alg() != jwt.SigningMethodRS256.Alg() {
			return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"])
		}
		return v.cfg.PublicKey, nil
	}

	opts := []jwt.ParserOption{jwt.WithValidMethods([]string{jwt.SigningMethodRS256.Alg()})}
	if iss := strings.TrimSpace(v.cfg.Issuer); iss != "" {
		opts = append(opts, jwt.WithIssuer(iss))
	}
	if aud := strings.TrimSpace(v.cfg.Audience); aud != "" {
		opts = append(opts, jwt.WithAudience(aud))
	}

	token, err := jwt.Parse(tokenString, keyFunc, opts...)
	if err != nil {
		return nil, fmt.Errorf("invalid jwt: %w", err)
	}
	if !token.Valid {
		return nil, errors.New("invalid jwt")
	}

	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		return nil, errors.New("invalid jwt claims")
	}

	p := &Principal{
		Sub:     claimString(claims, "sub"),
		AgentID: claimString(claims, "agent_id"),
		Roles:   claimStringSlice(claims, "roles"),
		Scopes:  claimStringSlice(claims, "scope", "scopes"),
	}
	if p.AgentID == "" && p.Sub != "" {
		p.AgentID = strings.TrimPrefix(p.Sub, "agent:")
	}
	if p.AgentID == "" {
		return nil, errors.New("jwt missing agent_id claim")
	}
	return p, nil
}

func claimString(claims jwt.MapClaims, keys ...string) string {
	for _, k := range keys {
		if v, ok := claims[k]; ok {
			switch t := v.(type) {
			case string:
				return strings.TrimSpace(t)
			}
		}
	}
	return ""
}

func claimStringSlice(claims jwt.MapClaims, keys ...string) []string {
	for _, k := range keys {
		v, ok := claims[k]
		if !ok {
			continue
		}
		switch t := v.(type) {
		case []any:
			out := make([]string, 0, len(t))
			for _, item := range t {
				if s, ok := item.(string); ok && s != "" {
					out = append(out, s)
				}
			}
			return out
		case []string:
			return t
		case string:
			if t != "" {
				return []string{t}
			}
		}
	}
	return nil
}
