package utils

import (
	"context"
	"crypto"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/patrickmn/go-cache"

	"github.com/tidbcloud/account/thirdparty/infra"
	"github.com/tidbcloud/tiinf/monitor/tilog"

	"go.uber.org/zap"

	"github.com/lestrrat-go/jwx/v2/jwa"
	"github.com/lestrrat-go/jwx/v2/jwk"
	"github.com/lestrrat-go/jwx/v2/jws"
	"github.com/lestrrat-go/jwx/v2/jwt"
	"github.com/pingcap/log"
)

var (
	dedicatedMap  = map[string]int{"length": 32, "prefixLen": 23, "base58StrLen": 8}
	serverlessMap = map[string]int{"length": 15, "prefixLen": 6, "base58StrLen": 8}
)

type Field struct {
	Name  string
	Value interface{}
}

type JWTContainer struct {
	Headers jws.Headers
	Tokens  jwt.Token
}
type GetDBUserReq struct {
	Email           string
	ClusterId       *uint64
	ServerlessToken *string
	IsDedicated     *bool
	TargetType      string
}

func Base64Encode(data []byte) string {
	return base64.RawURLEncoding.EncodeToString(data)
}

func Base64Decode(data string) ([]byte, error) {
	return base64.RawURLEncoding.DecodeString(data)
}

func NewContainer() *JWTContainer {
	c := &JWTContainer{
		Headers: jws.NewHeaders(),
		Tokens:  jwt.New(),
	}
	return c
}

func (c *JWTContainer) AddHeader(key string, value interface{}) error {
	return c.Headers.Set(key, value)
}

func (c *JWTContainer) AddToken(key string, value interface{}) error {
	return c.Tokens.Set(key, value)
}

func (c *JWTContainer) AddTokens(fields []Field) error {
	for _, f := range fields {
		if err := c.Tokens.Set(f.Name, f.Value); err != nil {
			return err
		}
	}

	return nil
}

func (c *JWTContainer) BuildSignatureMessage() ([]byte, error) {
	rawHeader, err := json.Marshal(c.Headers)
	if err != nil {
		return nil, err
	}

	rawPayload, err := json.Marshal(c.Tokens)
	if err != nil {
		return nil, err
	}

	header := Base64Encode(rawHeader)
	payload := Base64Encode(rawPayload)
	buf := make([]byte, len(header)+len(payload)+1)
	copy(buf, header)
	buf[len(header)] = '.'
	copy(buf[len(header)+1:], payload)
	return buf, nil
}

func (c *JWTContainer) BuildJWTTokenWithMessage(message []byte, signature []byte) string {
	encodedSignature := []byte(Base64Encode(signature))
	buf := make([]byte, len(message)+len(encodedSignature)+1)
	copy(buf, message)
	buf[len(message)] = '.'
	copy(buf[len(message)+1:], encodedSignature)
	return string(buf)
}

func (c *JWTContainer) BuildJWTToken(signature []byte) string {
	message, err := c.BuildSignatureMessage()
	if err != nil {
		return ""
	}
	return c.BuildJWTTokenWithMessage(message, signature)
}

func GetJwtToken(ctx context.Context, kms *KeyService, keyID string, email string, sub string, expiredTime time.Duration, iss string) (string, error) {

	container := NewContainer()
	if err := container.Headers.Set(jws.TypeKey, "JWT"); err != nil {
		return "", err
	}
	if err := container.Headers.Set(jws.AlgorithmKey, jwa.RS256); err != nil {
		return "", err
	}
	if err := container.Headers.Set(jws.KeyIDKey, keyID); err != nil {
		return "", err
	}

	payload := []Field{
		{"email", email},
		{jwt.SubjectKey, sub},
		{jwt.IssuedAtKey, time.Now().Unix()},
		{jwt.ExpirationKey, time.Now().Add(expiredTime).Unix()},
		{jwt.IssuerKey, iss},
	}

	if err := container.AddTokens(payload); err != nil {
		return "", err
	}

	messageToSign, err := container.BuildSignatureMessage()
	if err != nil {
		panic(err)
	}

	digest := Hash(crypto.SHA256, messageToSign)
	signature, err := kms.SignDigest(keyID, digest)
	if err != nil {
		panic(err)
	}
	token := container.BuildJWTTokenWithMessage(messageToSign, signature)
	tilog.Info(ctx, "jwt token", zap.Any("email", email))
	return token, nil
}

func VerifyJWTLocal(kms *KeyService, keyID string, jwtToken string) (bool, error) {

	content, _ := kms.ExportPublicKey(keyID)
	log.Info("ExportPublicKey", zap.Any("publickey", content))
	keySet := jwk.NewSet()

	key, err := DerToJWK(content)
	if err != nil {
		panic(err)
	}
	_ = key.Set(jwk.KeyIDKey, keyID)
	_ = key.Set(jwk.KeyUsageKey, jwk.ForSignature)
	_ = key.Set(jwk.AlgorithmKey, jwa.RS256)
	_ = keySet.AddKey(key)
	b, _ := json.MarshalIndent(keySet, "", "  ")
	//fmt.Printf("key: %s\n", b)
	log.Info("jsonKey", zap.Any("ke", b))

	fmt.Printf("JWT: %s\n", jwtToken)
	payload, err := jws.Verify([]byte(jwtToken), jws.WithKeySet(keySet))

	fmt.Println(string(payload))
	return err == nil, err
}

// GetDBUser
//
//	@Description: GetDBUser（SqlUser)
//	@Author wangletian 2022-12-28 22:36:28
//	@param ctx
//	@param req
//	@return *string
//	@return error
func GetDBUser(ctx context.Context, req GetDBUserReq) (*string, error) {
	tilog.Info(ctx, "GetDBUser req", zap.Any("req", req))
	var dbUser string
	var paramMap map[string]int
	if *req.IsDedicated {
		// dedicate cluster
		paramMap = dedicatedMap
		if strings.ToUpper(req.TargetType) == "APIKEY" {
			dbUser = calculateApiDBUser(ctx, req, paramMap)
		} else {
			dbUser = calculateDBUser(ctx, req, paramMap)
		}
	} else {
		// serverless cluster
		paramMap = serverlessMap
		if req.ServerlessToken == nil {
			c := infra.GetClusterCache(ctx)
			key := strconv.FormatUint(*req.ClusterId, 10) + "_serverless_token"
			value, found := c.Get(key)
			if found {
				str := fmt.Sprintf("%d", value)
				req.ServerlessToken = &str
				tilog.Info(ctx, "get ServerlessToken form cache", zap.Any("ClusterId", key), zap.Any("ServerlessToken", str))
			} else {

				cluster, err := infra.GetHandler().GetClusterById(ctx, *req.ClusterId, true)
				if err != nil {
					tilog.Info(ctx, "Call CM GetClusterById error", zap.Any("ClusterId", req.ClusterId))
					return nil, err
				}
				if cluster == nil {
					tilog.Info(ctx, "GetClusterById to  get ServerlessToken error , cluster is not exist ,ignore", zap.Any("clusterId", *req.ClusterId), zap.Any("dbUser", dbUser))
					return nil, errors.New(infra.CLUSTER_NOT_EXIST)
				}
				if !infra.IsSupportDBUser(ctx, cluster) {

					return nil, errors.New(infra.CLUSTER_NOT_SUPPORT_DBUSER)
				}

				if cluster.ServerlessToken == nil {
					return nil, errors.New("ServerlessToken is nil")
				}
				req.ServerlessToken = cluster.ServerlessToken
				c.Set(key, cluster.ServerlessToken, cache.DefaultExpiration)
				tilog.Info(ctx, "not found value form cache, set in cache", zap.Any("ClusterId", key), zap.Any("ServerlessToken", cluster.ServerlessToken))
			}
		}
		if strings.ToUpper(req.TargetType) == "APIKEY" {
			dbUser = calculateApiDBUser(ctx, req, paramMap)
		} else {
			dbUser = calculateDBUser(ctx, req, paramMap)
		}
		dbUser = *req.ServerlessToken + "." + dbUser
	}

	tilog.Info(ctx, "GetDBUser", zap.Any("dbUser", dbUser))
	return &dbUser, nil

}

func calculateApiDBUser(ctx context.Context, req GetDBUserReq, paramMap map[string]int) string {
	return req.Email + "_APIKEY"
}

// $email = replace($original_email, '@', '_')
// $db_account = prefix($email, 23)_prefix(base58(sha1($email)), 8)
func calculateDBUser(ctx context.Context, req GetDBUserReq, paramMap map[string]int) string {
	if len(req.Email) <= paramMap["length"] {
		tilog.Info(ctx, "calculate dbUser, dbUser = email", zap.Any("email", req.Email), zap.Any("dbUser", req.Email))
		return req.Email
	}
	//$email = replace($original_email, '@', '_')
	//$db_account = prefix($email, 23)_prefix(base58(sha1($email)), 8)

	//prefix($email, 23)
	newEmail := strings.Replace(req.Email, "@", "_", -1)
	rs := []rune(newEmail)
	subStr01 := string(rs[0:paramMap["prefixLen"]])

	// sha1($email)
	ins := sha1.New()
	ins.Write([]byte(newEmail))
	sha1Str := ins.Sum([]byte(""))

	// base58(sha1($email)
	base58Str := Base58Encoding(string(sha1Str))
	subStr02 := string([]rune(base58Str)[0:paramMap["base58StrLen"]])
	dbUser := subStr01 + "_" + subStr02
	tilog.Info(ctx, "calculate dbUser, dbUser != email", zap.Any("email", req.Email), zap.Any("dbUser", dbUser))
	return dbUser
}