package itoken

import (
	"context"
	"crypto/rsa"
	"encoding/base64"
	"fmt"
	"maps"
	"math/big"
	"slices"
	"time"

	jwt "github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"
	itkent "github.com/tidbcloud/devops/entity/itoken"
	"github.com/tidbcloud/devops/pkg/infra/config"
	"github.com/tidbcloud/devops/pkg/infra/consts"
	itksvcapi "github.com/tidbcloud/devops/pkg/service/itoken/api"

	"go.uber.org/dig"
)

var reservedClaims = []string{"iss", "aud", "sub", "iat", "exp", "jti"}

type ServiceParam struct {
	dig.In

	Conf *config.RootConfig
}

type svc struct {
	issuer     string
	defaultTTL time.Duration
	keyID      string

	privateKey *rsa.PrivateKey
	jwks       *itkent.JWKSet
}

func NewService(p ServiceParam) (itksvcapi.InternalTokenServiceAPI, error) {
	conf := p.Conf.Services.InternalToken
	privateKey, err := jwt.ParseRSAPrivateKeyFromPEM([]byte(conf.SigningKeyPEM))
	if err != nil {
		return nil, fmt.Errorf("parse jwt token signing key: %w", err)
	}

	return &svc{
		issuer:     consts.ServerAddr(),
		defaultTTL: *conf.DefaultTTL,
		keyID:      conf.KeyID,
		privateKey: privateKey,
		jwks:       buildJWKS(&privateKey.PublicKey, conf.KeyID),
	}, nil
}

func (s *svc) IssueJWTToken(_ context.Context, param itksvcapi.IssueJWTTokenParam) (*itkent.IssuedJWTToken, error) {
	if param.Subject == "" {
		return nil, fmt.Errorf("jwt token subject is required")
	}
	if param.Audience == "" {
		return nil, fmt.Errorf("jwt token audience is required")
	}
	for claim := range param.CustomClaims {
		if slices.Contains(reservedClaims, claim) {
			return nil, fmt.Errorf("custom claim %q is reserved", claim)
		}
	}

	now := time.Now()
	expireAt := now.Add(s.defaultTTL)
	if param.ExpiresAt != nil {
		expireAt = *param.ExpiresAt
	}
	if !expireAt.After(now) {
		return nil, fmt.Errorf("jwt token expiration must be in the future")
	}

	jti := uuid.NewString()
	if param.JWTID != nil {
		jti = *param.JWTID
	}

	claims := jwt.MapClaims{
		"iss": s.issuer,
		"aud": param.Audience,
		"sub": param.Subject,
		"iat": now.Unix(),
		"exp": expireAt.Unix(),
		"jti": jti,
	}
	maps.Copy(claims, param.CustomClaims)

	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	token.Header["kid"] = s.keyID

	signed, err := token.SignedString(s.privateKey)
	if err != nil {
		return nil, fmt.Errorf("sign jwt token: %w", err)
	}

	return &itkent.IssuedJWTToken{
		Token:    signed,
		ExpireAt: expireAt,
		Issuer:   s.issuer,
		Audience: param.Audience,
		KeyID:    s.keyID,
	}, nil
}

func (s *svc) GetJWKS(_ context.Context) (*itkent.JWKSet, error) {
	keys := make([]itkent.JWK, 0, len(s.jwks.Keys))
	keys = append(keys, s.jwks.Keys...)
	return &itkent.JWKSet{Keys: keys}, nil
}

func buildJWKS(publicKey *rsa.PublicKey, keyID string) *itkent.JWKSet {
	return &itkent.JWKSet{
		Keys: []itkent.JWK{
			{
				Kty: "RSA",
				Use: "sig",
				Alg: jwt.SigningMethodRS256.Alg(),
				Kid: keyID,
				N:   encodeBigInt(publicKey.N),
				E:   encodeBigInt(big.NewInt(int64(publicKey.E))),
			},
		},
	}
}

func encodeBigInt(i *big.Int) string {
	return base64.RawURLEncoding.EncodeToString(i.Bytes())
}