package jwtutil

import (
	"crypto/rsa"
	"encoding/base64"
	"errors"
	"fmt"
	"math/big"
	"os"
	"strings"
	"sync"
	"time"

	beego "github.com/beego/beego/v2/server/web"
	"github.com/golang-jwt/jwt/v5"
)

// 算法常量
const (
	AlgHS256 = "HS256"
	AlgRS256 = "RS256"
)

// DefaultKid 默认密钥 ID。历史 token 未携带 kid，统一按此 ID 处理。
const DefaultKid = "default"

// legacySecret 兼容用的内置密钥：仅用于解析历史已签发的 token。
// 生产环境必须在 app.conf 中配置 jwt_secret，否则会一直使用该固定值。
const legacySecret = "yunzer_jwt_secret_key"

// Claims JWT 载荷。旧字段保持不变，新增字段均为 omitempty，
// 历史 token 解析后为零值，不影响现有 78 处解析点。
type Claims struct {
	UserID   int    `json:"user_id"`
	Username string `json:"username"`
	TenantId int    `json:"tenant_id"` // 租户ID
	UserType string `json:"user_type"` // 用户类型："user" / "employee" / "platform" 等

	// ---- 统一认证中心（OIDC）扩展字段 ----
	ClientID string `json:"client_id,omitempty"` // 接入应用 client_id，即 aud 的业务标识
	Sid      string `json:"sid,omitempty"`       // 会话ID，用于单点登出/踢下线
	Scope    string `json:"scope,omitempty"`
	Amr      string `json:"amr,omitempty"` // 认证方式：pwd/sms/otp/wx/...

	// Events OIDC 事件声明，用于 Back-Channel Logout：
	// {"http://schemas.openid.net/event/backchannel-logout": {}}
	Events map[string]interface{} `json:"events,omitempty"`

	jwt.RegisteredClaims
}

// ---------------------------------------------------------------- 密钥管理

var (
	keyOnce   sync.Once
	hsKeys    map[string][]byte // kid -> HMAC 密钥
	rsaPriv   *rsa.PrivateKey   // RS256 签名
	rsaPub    *rsa.PublicKey    // RS256 验签
	rsaKid    string            // RS256 密钥 ID
	issuerVal string
)

func loadKeys() {
	hsKeys = map[string][]byte{DefaultKid: []byte(legacySecret)}

	// 主密钥（覆盖内置默认值）
	if s, _ := beego.AppConfig.String("jwt_secret"); strings.TrimSpace(s) != "" {
		hsKeys[DefaultKid] = []byte(strings.TrimSpace(s))
	}
	// 轮换密钥：jwt_secrets = kid1:secret1,kid2:secret2
	if s, _ := beego.AppConfig.String("jwt_secrets"); strings.TrimSpace(s) != "" {
		for _, item := range strings.Split(s, ",") {
			kv := strings.SplitN(strings.TrimSpace(item), ":", 2)
			if len(kv) == 2 && strings.TrimSpace(kv[0]) != "" && strings.TrimSpace(kv[1]) != "" {
				hsKeys[strings.TrimSpace(kv[0])] = []byte(strings.TrimSpace(kv[1]))
			}
		}
	}

	// RS256 密钥对：优先读文件路径，其次读内联 PEM
	privPEM := readKeyConf("jwt_rsa_private_key", "jwt_rsa_private_key_file")
	pubPEM := readKeyConf("jwt_rsa_public_key", "jwt_rsa_public_key_file")
	if privPEM != "" {
		if k, err := jwt.ParseRSAPrivateKeyFromPEM([]byte(privPEM)); err == nil {
			rsaPriv = k
			rsaPub = &k.PublicKey
		}
	}
	if pubPEM != "" {
		if k, err := jwt.ParseRSAPublicKeyFromPEM([]byte(pubPEM)); err == nil {
			rsaPub = k
		}
	}
	if rsaKid == "" {
		rsaKid, _ = beego.AppConfig.String("jwt_rsa_kid")
	}
	if rsaKid == "" {
		rsaKid = "rsa-1"
	}
	issuerVal, _ = beego.AppConfig.String("jwt_issuer")
}

func readKeyConf(inlineKey, fileKey string) string {
	if v, _ := beego.AppConfig.String(inlineKey); strings.TrimSpace(v) != "" {
		return strings.ReplaceAll(strings.TrimSpace(v), `\n`, "\n")
	}
	if p, _ := beego.AppConfig.String(fileKey); strings.TrimSpace(p) != "" {
		if b, err := os.ReadFile(strings.TrimSpace(p)); err == nil {
			return string(b)
		}
	}
	return ""
}

func ensureKeys() {
	keyOnce.Do(loadKeys)
}

// keyFunc 按 token header 的 alg + kid 选择验签密钥。
// 未携带 kid 或 kid 未注册时回落到默认密钥，保证历史 token 仍可解析。
func keyFunc(token *jwt.Token) (interface{}, error) {
	ensureKeys()
	kid, _ := token.Header["kid"].(string)
	switch token.Method.Alg() {
	case AlgHS256:
		if sec, ok := hsKeys[kid]; ok && kid != "" {
			return sec, nil
		}
		return hsKeys[DefaultKid], nil
	case AlgRS256:
		if rsaPub == nil {
			return nil, errors.New("服务端未配置 RS256 公钥")
		}
		return rsaPub, nil
	}
	return nil, fmt.Errorf("unsupported signing method: %v", token.Header["alg"])
}

// ---------------------------------------------------------------- 签发

// TokenOptions 认证中心签发参数
type TokenOptions struct {
	Alg      string // HS256 / RS256；留空时优先 RS256，未配置 RSA 则回落 HS256
	UserID   int    // 身份 ID（OIDC sub 用字符串形式）
	Username string
	TenantID int
	UserType string
	ClientID string
	Sid      string
	Scope    string
	Amr      string
	Subject  string
	Audience []string
	Jti      string                 // JWT ID，用于吊销（登出/踢下线）
	Events   map[string]interface{} // OIDC 事件声明（Back-Channel Logout 用）
	TTL      time.Duration          // 留空默认 30 分钟
	Kid      string
}

// SignToken 签发 token。RSA 未配置时自动回落 HS256，保证服务可启动。
func SignToken(opt TokenOptions) (string, error) {
	ensureKeys()

	alg := strings.ToUpper(strings.TrimSpace(opt.Alg))
	if alg == "" {
		alg = AlgRS256
	}
	if alg == AlgRS256 && rsaPriv == nil {
		alg = AlgHS256
	}

	ttl := opt.TTL
	if ttl <= 0 {
		ttl = 30 * time.Minute
	}
	now := time.Now()
	claims := &Claims{
		UserID:   opt.UserID,
		Username: opt.Username,
		TenantId: opt.TenantID,
		UserType: opt.UserType,
		ClientID: opt.ClientID,
		Sid:      opt.Sid,
		Scope:    opt.Scope,
		Amr:      opt.Amr,
		Events:   opt.Events,
		RegisteredClaims: jwt.RegisteredClaims{
			ID:        opt.Jti,
			Subject:   opt.Subject,
			Audience:  opt.Audience,
			ExpiresAt: jwt.NewNumericDate(now.Add(ttl)),
			IssuedAt:  jwt.NewNumericDate(now),
			NotBefore: jwt.NewNumericDate(now),
		},
	}
	if issuerVal != "" {
		claims.Issuer = issuerVal
	}
	if opt.Subject == "" && opt.UserID > 0 {
		claims.Subject = fmt.Sprintf("%d", opt.UserID)
	}

	var method jwt.SigningMethod
	switch alg {
	case AlgRS256:
		method = jwt.SigningMethodRS256
	default:
		method = jwt.SigningMethodHS256
	}
	token := jwt.NewWithClaims(method, claims)

	kid := strings.TrimSpace(opt.Kid)
	if kid == "" && alg == AlgRS256 {
		kid = rsaKid
	}
	if kid != "" {
		token.Header["kid"] = kid
	}

	var key interface{}
	if alg == AlgRS256 {
		key = rsaPriv
	} else {
		kidToUse := kid
		if kidToUse == "" {
			kidToUse = DefaultKid
		}
		if sec, ok := hsKeys[kidToUse]; ok {
			key = sec
		} else {
			key = hsKeys[DefaultKid]
		}
	}
	return token.SignedString(key)
}

// GenerateToken 生成 JWT token（兼容旧签名与行为）。
// 默认 HS256；若 app.conf 配置了 jwt_issuer 则带上 iss。
func GenerateToken(userID int, username string, tenantId int, userType string) (string, error) {
	ttl := 24 * time.Hour
	ensureKeys()
	now := time.Now()
	claims := &Claims{
		UserID:   userID,
		Username: username,
		TenantId: tenantId,
		UserType: userType,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(now.Add(ttl)),
			IssuedAt:  jwt.NewNumericDate(now),
			NotBefore: jwt.NewNumericDate(now),
		},
	}
	if issuerVal != "" {
		claims.Issuer = issuerVal
	}
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString(hsKeys[DefaultKid])
}

// ---------------------------------------------------------------- 解析

// ParseToken 解析并校验 JWT。自动识别 HS256 / RS256，函数签名与旧版一致。
func ParseToken(tokenString string) (*Claims, error) {
	claims := &Claims{}
	_, err := jwt.ParseWithClaims(tokenString, claims, keyFunc,
		jwt.WithValidMethods([]string{AlgHS256, AlgRS256}))
	if err != nil {
		return nil, err
	}
	return claims, nil
}

// ParseTokenRaw 解析 token 但不校验有效期，用于登出/吊销场景获取 jti。
func ParseTokenRaw(tokenString string) (*Claims, error) {
	claims := &Claims{}
	parser := jwt.NewParser(jwt.WithValidMethods([]string{AlgHS256, AlgRS256}), jwt.WithoutClaimsValidation())
	if _, err := parser.ParseWithClaims(tokenString, claims, keyFunc); err != nil {
		return nil, err
	}
	return claims, nil
}

// Issuer 返回配置的签发者（未配置时为空）
func Issuer() string {
	ensureKeys()
	return issuerVal
}

// JWK JSON Web Key（RS256 公钥）
type JWK struct {
	Kty string `json:"kty"`
	Use string `json:"use"`
	Alg string `json:"alg"`
	Kid string `json:"kid"`
	N   string `json:"n"`
	E   string `json:"e"`
}

// JWKS 返回 RSA 公钥集合（供 /auth/jwks.json 暴露）。
// 各应用本地用公钥验签即可，无需每次回调认证中心的 introspect 接口。
func JWKS() []JWK {
	ensureKeys()
	if rsaPub == nil || rsaPub.N == nil {
		return nil
	}
	return []JWK{{
		Kty: "RSA",
		Use: "sig",
		Alg: AlgRS256,
		Kid: rsaKid,
		N:   base64.RawURLEncoding.EncodeToString(rsaPub.N.Bytes()),
		E:   base64.RawURLEncoding.EncodeToString(big.NewInt(int64(rsaPub.E)).Bytes()),
	}}
}

// HasRSA 是否已配置 RS256 密钥对（决定认证中心能否签发非对称 token）
func HasRSA() bool {
	ensureKeys()
	return rsaPriv != nil && rsaPub != nil
}

// Kid 返回当前 RS256 密钥 ID
func Kid() string {
	ensureKeys()
	return rsaKid
}