package oauth2

import (
	"context"
	"fmt"
	"time"

	"go.uber.org/zap"
	"gorm.io/gorm"

	"github.com/prometheus/client_golang/prometheus"
	"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/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 (
	deviceCodeExpiresInSeconds = 600
	pollIntervalSeconds        = 5

	// 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.
	deviceCodeEncryptSalt        = ""
	deviceCodeEncryptInterations = 20000
)

// https://datatracker.ietf.org/doc/html/rfc8628#section-3.1
func (o *Oauth2ServiceImpl) HandleDeviceAuthorizationRequest(ctx context.Context, clientID, clientIP, userAgent string) (*dto.DeviceAuthorizationResponse, error) {
	db := o.db.WithContext(ctx)
	app, err := repository.GetOauth2ApplicationByClientID(ctx, db, clientID)
	if err != nil {
		return nil, errs.ErrOauth2InvalidClient
	}
	if app == nil || app.GrantType != entity.Oauth2DeviceCode {
		return nil, errs.ErrOauth2InvalidClient
	}

	verificationURI := config.Global().Oauth2Config.VerificationURI

	userCode, deviceCode, err := createDeviceGrantWithRetry(ctx, db, &entity.Oauth2DeviceGrant{
		ApplicationID: uint64(app.ID),
		ExpiresIn:     deviceCodeExpiresInSeconds,
		// Scopes and ProjectIDS will be filled after user authorization
		Scopes:     "[]",
		ProjectIDs: entity.ProjectIDs{IDs: []uint64{}},
		ClientIP:   clientIP,
		UserAgent:  userAgent,
	})
	if err != nil {
		tilog.Error(ctx, "failed to create oauth2 device grant", zap.Error(err))
		return nil, errs.ErrOauth2ServerError
	}

	return &dto.DeviceAuthorizationResponse{
		DeviceCode:              deviceCode,
		UserCode:                userCode,
		VerificationURI:         verificationURI,
		VerificationURIComplete: verificationURI + "?user_code=" + userCode,
		ExpiresIn:               deviceCodeExpiresInSeconds,
		Interval:                pollIntervalSeconds,
	}, nil
}

// https://datatracker.ietf.org/doc/html/rfc8628#section-3.4
func (o *Oauth2ServiceImpl) HandleDeviceTokenRequest(ctx context.Context, clientID, deviceCode string) (*dto.TokenResponse, error) {
	var (
		resp     *dto.TokenResponse
		oauthErr error
	)

	encrptedDeviceCode := crypto.NewPbdkf2Sha512Encryptor().Encrypt(
		deviceCode, deviceCodeEncryptSalt, deviceCodeEncryptInterations,
	)

	db := o.db.WithContext(ctx)
	app, dbErr := repository.GetOauth2ApplicationByClientID(ctx, db, clientID)
	if dbErr != nil || app == nil || app.GrantType != entity.Oauth2DeviceCode {
		return nil, errs.ErrOauth2InvalidClient
	}

	err := db.Transaction(func(tx *gorm.DB) error {
		grant, err := repository.GetOauth2DeviceGrantByDeviceCodeWithLock(ctx, tx, encrptedDeviceCode)
		if err != nil {
			oauthErr = errs.ErrOauth2IncorrectDeviceCode
			return nil
		}

		if grant.ApplicationID != uint64(app.ID) {
			oauthErr = errs.ErrOauth2IncorrectDeviceCode
			return nil
		}

		now := time.Now()

		switch grant.Status {
		case entity.Oauth2GrantStatusDeny:
			oauthErr = errs.ErrOauth2AccessDenied
			return nil
		case entity.Oauth2GrantStatusPending:
			if deviceGrantRevokedOrExpired(grant, now) {
				tilog.Info(ctx, "The device_code has expired or revoked")
				oauthErr = errs.ErrOauth2ExpiredToken
				return nil
			}

			slowDown := false
			if grant.LastPolledAt != nil {
				nextPollTime := grant.LastPolledAt.Add(time.Duration(pollIntervalSeconds) * time.Second)
				tilog.Info(ctx, "next poll time", zap.Time("nextPollTime", nextPollTime))
				if now.Before(nextPollTime) {
					slowDown = true
				}
			}
			if slowDown {
				oauthErr = errs.ErrOauth2SlowDown
				tilog.Info(ctx, "Calling token api too frequently", zap.Any("response err", oauthErr.Error()))
				return nil
			}

			grant.LastPolledAt = &now
			err = repository.UpdateOauth2DeviceGrant(
				ctx, tx, grant,
				[]entity.Oauth2DeviceGrantDBSchemaField{
					entity.Oauth2DeviceGrantDBSchema.LastPolledAt,
				},
			)
			if err != nil {
				tilog.Error(ctx, "failed to update oauth2 device grant", zap.Error(err))
				return err
			}
			oauthErr = errs.ErrOauth2AuthorizationPending
			return nil
		case entity.Oauth2GrantStatusApprove:
			if deviceGrantRevokedOrExpired(grant, now) {
				tilog.Info(ctx, "The device_code has expired or revoked")
				oauthErr = errs.ErrOauth2ExpiredToken
				return nil
			}

			akObj, ak, rk, err := createAccessTokenByDeviceGrant(ctx, tx, grant, app, o.accountProviderClient)
			if err != nil {
				tilog.Error(ctx, "failed to create access token", zap.Error(err))
				return err
			}
			grant.RevokedAt = &now
			err = repository.UpdateOauth2DeviceGrant(
				ctx, tx, grant,
				[]entity.Oauth2DeviceGrantDBSchemaField{
					entity.Oauth2DeviceGrantDBSchema.RevokedAt,
				},
			)
			if err != nil {
				tilog.Error(ctx, "failed to update oauth2 device grant", zap.Error(err))
				return err
			}
			resp = &dto.TokenResponse{
				AccessToken:  ak,
				TokenType:    TokenTypeBearer,
				ExpiresIn:    akObj.AccessTokenExpiresIn,
				RefreshToken: rk,
			}
			return nil
		default:
			return fmt.Errorf("got invalid oauth2 device grant status %d", grant.Status)
		}
	})

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

func deviceGrantRevokedOrExpired(grant *entity.Oauth2DeviceGrant, now time.Time) bool {
	if grant.RevokedAt != nil {
		return true
	}

	expiresAt := grant.CreatedAt.Add(time.Duration(grant.ExpiresIn) * time.Second)
	return now.After(expiresAt)
}

func createAccessTokenByDeviceGrant(
	ctx context.Context, db *gorm.DB, grant *entity.Oauth2DeviceGrant, application *entity.Oauth2Application, accountProvider account.AccountAPIServiceClient,
) (obj *entity.Oauth2AccessToken, accessToken string, refreshToken string, err error) {
	obj = &entity.Oauth2AccessToken{
		AccessTokenExpiresIn:  application.AccessTokenExpiresIn,
		RefreshTokenExpiresIn: application.RefreshTokenExpiresIn,
		ApplicationID:         grant.ApplicationID,
		UserID:                grant.UserID,
		UserEmail:             grant.UserEmail,
		OrgID:                 grant.OrgID,
		Scopes:                entity.Oauth2Scope{},
		ProjectIDs:            grant.ProjectIDs,
		GrantAllProjects:      grant.GrantAllProjects,
		GrantType:             entity.Oauth2DeviceCode,
	}
	accessToken, refreshToken, err = createAccessTokenWithRetry(ctx, db, obj, accountProvider)
	return
}

// createDeviceGrantWithRetry will retry to create device grant if the generated userCode/deviceCode is duplicated
func createDeviceGrantWithRetry(ctx context.Context, db *gorm.DB, obj *entity.Oauth2DeviceGrant) (userCode, deviceCode string, err error) {
	for i := 0; i < maxDBDuplicateRetryTimes; i++ {
		userCode = NewUserCode()
		deviceCode = NewDeviceCode()
		encrptedDeviceCode := crypto.NewPbdkf2Sha512Encryptor().Encrypt(
			deviceCode, deviceCodeEncryptSalt, deviceCodeEncryptInterations,
		)
		obj.UserCode = userCode
		obj.DeviceCode = encrptedDeviceCode
		err = repository.CreateOauth2DeviceGrant(ctx, db, obj)
		if err == nil {
			return userCode, deviceCode, nil
		} else if errs.IsDuplicateEntryError(err) {
			metrics.DBDuplicateEntryCount.With(prometheus.Labels{
				"type": "oauth2_device_grant",
			}).Inc()
			continue
		} else {
			// if not duplicated key error, return withouth retry
			return "", "", err
		}
	}
	return "", "", fmt.Errorf("Failed to create oauth2 device grant after %d retries", maxDBDuplicateRetryTimes)
}