package oauth2

import (
	"context"
	"fmt"
	"github.com/duke-git/lancet/v2/slice"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/tidbcloud/auth-server/thirdparty"
	"go.uber.org/zap"
	"gorm.io/gorm"
	"regexp"
	"time"

	"github.com/tidbcloud/account/idl/pbgen/proto/account"
	"github.com/tidbcloud/auth-server/commons/config"
	"github.com/tidbcloud/auth-server/commons/crypto"
	"github.com/tidbcloud/auth-server/commons/errs"
	"github.com/tidbcloud/auth-server/commons/grpc"
	"github.com/tidbcloud/auth-server/commons/metrics"
	dto "github.com/tidbcloud/auth-server/internal/dto/oauth2"
	"github.com/tidbcloud/auth-server/internal/entity"
	"github.com/tidbcloud/auth-server/internal/repository"
	"github.com/tidbcloud/tiinf/monitor/tilog"
)

const (
	// An empty salt is used because we need to look tokens up solely by
	// their hashed value. Additionally, tokens are always cryptographically
	// pseudo-random and unique, therefore salting provides no
	// additional security.
	tokenEncryptSalt        = ""
	tokenEncryptInterations = 20000
)

var pingcapTms = []string{"TMS", "TLMS"}

// HandleDeviceAuthorizationRequest handles refresh token request.
// https://datatracker.ietf.org/doc/html/rfc6749#section-6
func (o *Oauth2ServiceImpl) HandleRefreshTokenRequest(ctx context.Context, clientID, clientSecret, refreshToken string) (*dto.TokenResponse, error) {
	var (
		resp     *dto.TokenResponse
		oauthErr error
	)

	db := o.db.WithContext(ctx)
	app, dbErr := repository.GetOauth2ApplicationByClientID(ctx, db, clientID)
	if dbErr != nil {
		tilog.Info(ctx, "req client_id is incorrect", zap.Any("clientID", clientID))
		return nil, errs.ErrOauth2InvalidClient
	}
	passwordEncryptor := crypto.NewPbdkf2Sha512Encryptor()
	if !passwordEncryptor.Compare(clientSecret, app.ClientSecret) {
		tilog.Info(ctx, "req client_secret is incorrect", zap.Any("clientSecret", clientSecret))
		return nil, errs.ErrOauth2InvalidClient
	}
	encryptedRefreshToken := passwordEncryptor.Encrypt(
		refreshToken, tokenEncryptSalt, tokenEncryptInterations,
	)
	// use transaction to lock the row, avoid concurrent update
	err := db.Transaction(func(tx *gorm.DB) error {
		akObj, err := repository.GetOauth2AccessTokenByRefreshTokenWithLock(ctx, tx, encryptedRefreshToken)
		if err != nil {
			// ignore this error, it is typically caused by invalid refreshToken or not found
			tilog.Info(ctx, "GetOauth2AccessTokenByRefreshToke", zap.Any("err", err.Error()))
			oauthErr = errs.ErrOauth2InvalidGrant
			return nil
		}
		if akObj.ApplicationID != uint64(app.ID) || isRefreshTokenExpired(akObj) {
			oauthErr = errs.ErrOauth2InvalidGrant
			return nil
		}

		newAkObj, accessToken, refreshToken, err := createAccessTokenByRefresh(ctx, tx, akObj, o.accountProviderClient)
		if err != nil {
			tilog.Error(ctx, "failed to create access token by refresh", zap.Error(err))
			return err
		}

		resp = &dto.TokenResponse{
			AccessToken:  accessToken,
			TokenType:    TokenTypeBearer,
			ExpiresIn:    newAkObj.AccessTokenExpiresIn,
			RefreshToken: refreshToken,
		}
		return nil
	})

	if err != nil {
		tilog.Error(ctx, "failed to handle refresh token request", zap.Error(err))
		return nil, errs.ErrOauth2ServerError
	}
	return resp, oauthErr
}

// HandleRevokeTokenRequest handles revoke token request.
// Do not need to return any response, hide the details from the client.
// https://datatracker.ietf.org/doc/html/rfc7009
func (o *Oauth2ServiceImpl) HandleRevokeTokenRequest(ctx context.Context, clientID, clientSecret, token, tokenTypeHint string) {
	db := o.db.WithContext(ctx)
	app, err := repository.GetOauth2ApplicationByClientID(ctx, db, clientID)
	if err != nil {
		// ignore this error, it is typically caused by invalid clientID
		tilog.Debug(ctx, "failed to get application by clientID", zap.Error(err), zap.String("clientID", clientID))
		return
	}
	passwordEncryptor := crypto.NewPbdkf2Sha512Encryptor()
	if !passwordEncryptor.Compare(clientSecret, app.ClientSecret) {
		tilog.Debug(ctx, "client secret not match", zap.String("clientID", clientID))
		return
	}
	encryptedToken := passwordEncryptor.Encrypt(token, tokenEncryptSalt, tokenEncryptInterations)

	switch tokenTypeHint {
	// default is access_token
	case "access_token", "":
		err = revokeByAccessToken(ctx, db, encryptedToken, uint64(app.ID))
	case "refresh_token":
		err = revokeByRefreshToken(ctx, db, encryptedToken, uint64(app.ID))
	default:
		return
	}
	if err != nil {
		tilog.Error(
			ctx, "failed to revoke token", zap.Error(err), zap.String("clientID", clientID),
			zap.String("token", token), zap.String("tokenTypeHint", tokenTypeHint),
		)
	}
}

func (o *Oauth2ServiceImpl) RevokeTokenByFilter(ctx context.Context, req dto.RevokeTokenFilter) error {
	db := o.db.WithContext(ctx)

	return repository.RevokeOauth2AccessTokenByFilter(ctx, db, req)
}

func (o *Oauth2ServiceImpl) GetUserInfo(ctx context.Context, rawAccessToken string) (*dto.UserInfo, error) {

	accessToken, err := o.GetAccessTokenIdentity(ctx, rawAccessToken)
	if err != nil {
		return nil, errs.ErrOauth2InvalidAccessToken
	}
	if accessToken != nil {
		accountClient := grpc.GetAccountProviderSvcClient()

		req := &account.GetUserByIDRequest{

			UserId: accessToken.UserID,
		}
		response, err := accountClient.GetUserByID(ctx, req)
		if err != nil {
			tilog.Error(ctx, "ListUsersByFilter fail", zap.Error(err))
			return nil, err
		}
		if response != nil {
			userinfo := &dto.UserInfo{
				Email:     accessToken.UserEmail,
				FirstName: response.User.FirstName,
				LastName:  response.User.LastName,
				Username:  response.User.FirstName + " " + response.User.LastName,
			}
			// if tms app  , check pingcap email
			app, err := o.GetApplicationByID(ctx, accessToken.ApplicationID)
			if err != nil {
				tilog.Error(ctx, "GetApplicationByID fail", zap.Error(err))
				return nil, err
			}
			// tms return  tms userinfo
			if slice.Contain(pingcapTms, app.ClientID) {
				err = setTmsUserInfo(ctx, accessToken.UserEmail, accessToken.OrgID, userinfo)

			}

			tilog.Info(ctx, "GetUserInfoByAccessToken  result", zap.Any("userinfo", userinfo), zap.Any("err", err))
			return userinfo, err
		}

	}
	tilog.Info(ctx, "GetUserInfoByAccessToken  cant found the accessToken  in the database")
	return nil, nil
}

func setTmsUserInfo(ctx context.Context, email string, orgId uint64, userinfo *dto.UserInfo) error {
	// pingcap email is not allowed login to tms
	matched, err := regexp.MatchString(`^[a-zA-Z0-9._%+-]+@(pingcap.com)|(pingcap.cn)$`, email)
	if err != nil {
		tilog.Error(ctx, "regexp.MatchString fail", zap.Error(err))
	}
	if matched {
		tilog.Info(ctx, "pingcap email is not allowed", zap.String("email", email))
		return errs.ErrOauth2AccessDenied
	}
	tmsUser, err := thirdparty.GetTmsUsername(ctx, email)
	if err != nil {
		tilog.Error(ctx, "GetTmsUsername fail", zap.Error(err))
		return err
	}
	if tmsUser != nil {
		userinfo.Username = tmsUser.Login
		userinfo.FirstName = tmsUser.FirstName
		userinfo.LastName = tmsUser.LastName
		if userinfo.LastName == "" {
			tilog.Info(ctx, "set lastName =  .", zap.Any(" accessToken.UserEmail", email))
			userinfo.LastName = "."
		}
		tilog.Info(ctx, "set username =tmsUsername", zap.Any("tmsUser", tmsUser))

	} else {
		tilog.Info(ctx, "user not exits tms ,set ", zap.Any("tmsUser", tmsUser))
		accountClient := grpc.GetAccountProviderSvcClient()

		req := &account.GetOrgByIDRequest{
			Id: orgId,
		}
		response, err := accountClient.GetOrgByID(ctx, req)
		if err != nil {
			tilog.Error(ctx, "GetOrgByID fail", zap.Error(err))
			return err
		}
		if response != nil && response.Org != nil {
			userinfo.CompanyName = response.Org.Name
			userinfo.Country = response.Org.Country

			tilog.Info(ctx, "set company_name ,country from org", zap.Any("tmsUser", tmsUser))

		}
	}
	return nil
}
func (o *Oauth2ServiceImpl) GetAccessTokenInfo(ctx context.Context, rawAccessToken string) (*dto.AccessTokenInfo, error) {

	accessToken, err := o.GetAccessTokenIdentity(ctx, rawAccessToken)
	if err != nil {
		return nil, errs.ErrOauth2InvalidAccessToken
	}
	if accessToken != nil {

		accessTokenInfo := &dto.AccessTokenInfo{
			UserEmail: accessToken.UserEmail,
			UserId:    accessToken.UserID,
			OrgId:     accessToken.OrgID,
			GrantType: string(accessToken.GrantType),

			ProjectIds:       accessToken.ProjectIDs.IDs,
			GrantAllProjects: accessToken.GrantAllProjects,
			Scopes:           accessToken.Scopes,
		}
		tilog.Info(ctx, "GetAccessTokenInfo success", zap.Any("userinfo", accessTokenInfo))
		return accessTokenInfo, nil

	}
	tilog.Info(ctx, "GetUserInfoByAccessToken  cant found the accessToken  in the database")
	return nil, nil
}

func (o *Oauth2ServiceImpl) GetUserLastAccessToken(ctx context.Context, req dto.TokenFilter) (entity.Oauth2AccessToken, error) {
	db := o.db.WithContext(ctx)
	ak, err := repository.GetUserLastAccessToken(ctx, db, req)
	if err != nil {
		return ak, err
	}
	// not expired
	if !isAccessTokenExpired(&ak) {
		return ak, nil
	}
	tilog.Info(ctx, "UserLastAccessToken is expired ", zap.Any("AccessTokenId", ak.ID), zap.Any("CreatedAt", ak.CreatedAt), zap.Any("AccessTokenExpiresIn", ak.AccessTokenExpiresIn))
	return entity.Oauth2AccessToken{}, nil
}

func revokeByAccessToken(ctx context.Context, db *gorm.DB, accessToken string, applicationID uint64) error {
	akObj, err := repository.GetOauth2AccessTokenByAccessToken(ctx, db, accessToken)
	if err != nil {
		// ignore this error, it is typically caused by invalid accessToken or not found
		return nil
	}
	// if the token is already revoked or not belong to the application, do nothing
	if akObj.RevokedAt != nil || akObj.ApplicationID != applicationID {
		return nil
	}
	return repository.RevokeOauth2AccessToken(ctx, db, akObj.ID)
}

func revokeByRefreshToken(ctx context.Context, db *gorm.DB, refreshToken string, applicationID uint64) error {
	akObj, err := repository.GetOauth2AccessTokenByRefreshToken(ctx, db, refreshToken)
	if err != nil {
		// ignore this error, it is typically caused by invalid refreshToken or not found
		return nil
	}
	// if the token is already revoked or not belong to the application, do nothing
	if akObj.RevokedAt != nil || akObj.ApplicationID != applicationID {
		return nil
	}
	return repository.RevokeOauth2AccessToken(ctx, db, akObj.ID)
}

func isAccessTokenExpired(akObj *entity.Oauth2AccessToken) bool {
	if akObj.RevokedAt != nil {
		return true
	}
	// 0 means never expired
	if akObj.AccessTokenExpiresIn == 0 {
		return false
	}
	return akObj.CreatedAt.Add(time.Duration(akObj.AccessTokenExpiresIn) * time.Second).Before(time.Now())
}

func isRefreshTokenExpired(akObj *entity.Oauth2AccessToken) bool {
	if akObj.RevokedAt != nil {
		return true
	}
	// 0 means never expired
	if akObj.RefreshTokenExpiresIn == 0 {
		return false
	}
	return akObj.CreatedAt.Add(time.Duration(akObj.RefreshTokenExpiresIn) * time.Second).Before(time.Now())
}

func createAccessTokenByRefresh(ctx context.Context, db *gorm.DB, oldAkObj *entity.Oauth2AccessToken, accountProvider account.AccountAPIServiceClient) (obj *entity.Oauth2AccessToken, accessToken string, refreshToken string, err error) {
	now := time.Now()
	// revoke old access token
	oldAkObj.RevokedAt = &now
	err = repository.UpdateOauth2AccessToken(
		ctx, db, oldAkObj,
		[]entity.Oauth2AccessTokenDBSchemaField{
			entity.Oauth2AccessTokenDBSchema.RevokedAt,
		},
	)
	if err != nil {
		tilog.Error(ctx, "failed to revoke old access token", zap.Error(err))
		return nil, "", "", err
	}

	// create new access token
	obj = &entity.Oauth2AccessToken{
		AccessTokenExpiresIn:  oldAkObj.AccessTokenExpiresIn,
		RefreshTokenExpiresIn: oldAkObj.RefreshTokenExpiresIn,
		ApplicationID:         oldAkObj.ApplicationID,
		UserID:                oldAkObj.UserID,
		UserEmail:             oldAkObj.UserEmail,
		OrgID:                 oldAkObj.OrgID,
		Scopes:                oldAkObj.Scopes,
		ProjectIDs:            oldAkObj.ProjectIDs,
		GrantAllProjects:      oldAkObj.GrantAllProjects,
		GrantType:             oldAkObj.GrantType,
	}
	accessToken, refreshToken, err = createAccessTokenWithRetry(ctx, db, obj, accountProvider)
	return
}

// createAccessTokenWithRetry will retry to create access token if the generated accessToken/refreshToken is duplicated
func createAccessTokenWithRetry(ctx context.Context, db *gorm.DB, obj *entity.Oauth2AccessToken, accountProvider account.AccountAPIServiceClient) (accessToken string, refreshToken string, err error) {
	for i := 0; i < maxDBDuplicateRetryTimes; i++ {
		accessToken = NewAccessToken()
		refreshToken = NewRefreshToken()

		encryptor := crypto.NewPbdkf2Sha512Encryptor()
		encrptedAccessToken := encryptor.Encrypt(
			accessToken, tokenEncryptSalt, tokenEncryptInterations,
		)
		encrptedRefreshToken := encryptor.Encrypt(
			refreshToken, tokenEncryptSalt, tokenEncryptInterations,
		)

		obj.AccessToken = encrptedAccessToken
		obj.RefreshToken = encrptedRefreshToken

		err = repository.CreateOauth2AccessToken(ctx, db, obj)
		if err == nil {
			// SaveAuthorizationData to Account
			err = syncAccessTokenToAccount(ctx, obj, accountProvider)
			if err != nil {
				tilog.Error(ctx, "failed to sync access token to account", zap.Error(err))
				return "", "", err
			}
			return accessToken, refreshToken, nil
		} else if errs.IsDuplicateEntryError(err) {
			metrics.DBDuplicateEntryCount.With(prometheus.Labels{
				"type": "oauth2_access_token",
			}).Inc()
			continue
		} else {
			// if not duplicated key error, return withouth retry
			return "", "", err
		}
	}
	return "", "", fmt.Errorf("Failed to create oauth2 access token after %d retries", maxDBDuplicateRetryTimes)
}

// SaveAuthorizationData to Account
func syncAccessTokenToAccount(ctx context.Context, accessToken *entity.Oauth2AccessToken, accountProvider account.AccountAPIServiceClient) error {
	// only sync when it's a code grant and not local env
	if config.IsLocalEnv() || accessToken.GrantType != entity.Oauth2AuthorizationCode {
		return nil
	}

	resp, err := accountProvider.SaveAuthorizationData(ctx, &account.SaveAuthorizationDataReq{
		UserId:        accessToken.UserID,
		AccessTokenId: uint64(accessToken.ID),
		OrgId:         accessToken.OrgID,
		ProjectIds:    accessToken.ProjectIDs.IDs,
		OrgRole:       accessToken.Scopes.Org,
		ProjectRole:   accessToken.Scopes.Project,
	})
	if err = grpc.FormatError(ctx, err, resp.BaseResp); err != nil {
		tilog.Error(ctx, "failed to call accountProviderClient.SaveAuthorizationData", zap.Error(err))
		return err
	}
	return nil
}