package jwt_manager

import (
	"errors"
	"fmt"
	"github.com/golang-jwt/jwt/v5"
	"mingyang-admin-app-rpc/internal/config"
	"time"
)

type JWTManager struct {
	accessTokenSecret  []byte
	refreshTokenSecret []byte
	accessTokenExpiry  time.Duration
	refreshTokenExpiry time.Duration
	issuer             string
}

type Claims struct {
	UserID uint64 `json:"user_id"`
	Type   string `json:"type"` // "access" or "refresh"
	jwt.RegisteredClaims
}

func NewJWTManager(config *config.JWTConfig) *JWTManager {
	return &JWTManager{
		accessTokenSecret:  []byte(config.AccessTokenSecret),
		refreshTokenSecret: []byte(config.RefreshTokenSecret),
		accessTokenExpiry:  config.AccessTokenExpiry,
		refreshTokenExpiry: config.RefreshTokenExpiry,
		issuer:             config.Issuer,
	}
}

// GenerateAccessToken 生成访问令牌
func (m *JWTManager) GenerateAccessToken(userID uint64) (string, *Claims, error) {
	return m.generateToken(userID, "access", m.accessTokenExpiry, m.accessTokenSecret)
}

// GenerateRefreshToken 生成刷新令牌
func (m *JWTManager) GenerateRefreshToken(userID uint64) (string, *Claims, error) {
	return m.generateToken(userID, "refresh", m.refreshTokenExpiry, m.refreshTokenSecret)
}

func (m *JWTManager) generateToken(userID uint64, tokenType string, expiry time.Duration, secret []byte) (string, *Claims, error) {
	claims := &Claims{
		UserID: userID,
		Type:   tokenType,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(expiry)),
			IssuedAt:  jwt.NewNumericDate(time.Now()),
			Issuer:    m.issuer,
			Subject:   fmt.Sprint(userID),
		},
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	tokenString, err := token.SignedString(secret)
	if err != nil {
		return "", nil, err
	}

	return tokenString, claims, nil
}

// VerifyAccessToken 验证访问令牌
func (m *JWTManager) VerifyAccessToken(tokenString string) (*Claims, error) {
	return m.verifyToken(tokenString, m.accessTokenSecret, "access")
}

// VerifyRefreshToken 验证刷新令牌
func (m *JWTManager) VerifyRefreshToken(tokenString string) (*Claims, error) {
	return m.verifyToken(tokenString, m.refreshTokenSecret, "refresh")
}

func (m *JWTManager) verifyToken(tokenString string, secret []byte, expectedType string) (*Claims, error) {
	token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return secret, nil
	})

	if err != nil {
		return nil, err
	}

	if claims, ok := token.Claims.(*Claims); ok && token.Valid {
		if claims.Type != expectedType {
			return nil, errors.New("invalid token type")
		}
		return claims, nil
	}

	return nil, errors.New("invalid token")
}

// GenerateTokenPair 生成令牌对
func (m *JWTManager) GenerateTokenPair(userID uint64) (*TokenPair, error) {
	accessToken, accessClaims, err := m.GenerateAccessToken(userID)
	if err != nil {
		return nil, err
	}

	refreshToken, refreshClaims, err := m.GenerateRefreshToken(userID)
	if err != nil {
		return nil, err
	}

	return &TokenPair{
		AccessToken:           accessToken,
		RefreshToken:          refreshToken,
		AccessTokenExpiresAt:  accessClaims.ExpiresAt.Time,
		RefreshTokenExpiresAt: refreshClaims.ExpiresAt.Time,
		TokenType:             "Bearer",
	}, nil
}

type TokenPair struct {
	AccessToken           string    `json:"access_token"`
	RefreshToken          string    `json:"refresh_token"`
	AccessTokenExpiresAt  time.Time `json:"access_token_expires_at"`
	RefreshTokenExpiresAt time.Time `json:"refresh_token_expires_at"`
	TokenType             string    `json:"token_type"`
}