package oauth2_test

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"

	"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/internal/entity"
)

type HandleDeviceAuthorizationSuite struct {
	Oauth2Suite
}

func TestHandleDeviceAuthorizationSuite(t *testing.T) {
	suite.Run(t, new(HandleDeviceAuthorizationSuite))
}

func (s *HandleDeviceAuthorizationSuite) TestHandleDeviceAuthorizationRequest() {
	// invalid cases
	invalidCases := []struct {
		caseName  string
		clientID  string
		clientIP  string
		userAgent string
	}{
		{"empty", "", "", ""},
		{"empty2", "", "invalid", ""},
		{"not exist", "invalid", "", ""},
		{"invalid grant type", s.authorizationCodeApp.ClientID, "127.0.0.1", "ticloud"},
	}
	for _, c := range invalidCases {
		s.T().Run(c.caseName, func(t *testing.T) {
			resp, err := s.oauth2Service.HandleDeviceAuthorizationRequest(s.ctx, c.clientID, c.clientIP, c.userAgent)
			assert.NotNil(t, err)
			assert.Nil(t, resp)

			count, dbErr := entity.NewOauth2DeviceGrantQuerySet(s.db).Count()
			assert.Equal(t, int64(0), count)
			assert.Nil(t, dbErr)
		})
	}

	// valid cases
	resp, err := s.oauth2Service.HandleDeviceAuthorizationRequest(s.ctx, s.deviceFlowApp.ClientID, "127.0.0.1", "ticloud")
	assert.Nil(s.T(), err)
	require.NotNil(s.T(), resp)
	assert.Equal(s.T(), resp.VerificationURI, config.Global().Oauth2Config.VerificationURI)
	assert.Equal(s.T(), resp.VerificationURIComplete, config.Global().Oauth2Config.VerificationURI+"?user_code="+resp.UserCode)
	assert.Equal(s.T(), len(resp.UserCode), 9)
	assert.True(s.T(), len(resp.DeviceCode) > 20)

	count, dbErr := entity.NewOauth2DeviceGrantQuerySet(s.db).Count()
	require.Nil(s.T(), dbErr)
	assert.Equal(s.T(), int64(1), count)
	var grantObj entity.Oauth2DeviceGrant
	dbErr = entity.NewOauth2DeviceGrantQuerySet(s.db).One(&grantObj)
	require.Nil(s.T(), dbErr)
	assert.NotEqual(s.T(), resp.DeviceCode, grantObj.DeviceCode)
	encrptedDeviceCode := crypto.NewPbdkf2Sha512Encryptor().Encrypt(
		resp.DeviceCode, "", 20000,
	)
	assert.Equal(s.T(), encrptedDeviceCode, grantObj.DeviceCode)
	assert.Equal(s.T(), resp.UserCode, grantObj.UserCode)
	assert.Equal(s.T(), resp.ExpiresIn, grantObj.ExpiresIn)
	assert.Equal(s.T(), resp.Interval, 5)
	assert.Equal(s.T(), grantObj.Status, entity.Oauth2GrantStatusPending)
}

type HandleDeviceTokenSuite struct {
	Oauth2Suite
}

func TestHandleDeviceTokenSuite(t *testing.T) {
	suite.Run(t, new(HandleDeviceTokenSuite))
}

func (s *HandleDeviceTokenSuite) TestHandleDeviceTokenRequestInvalid() {
	resp, err := s.oauth2Service.HandleDeviceAuthorizationRequest(s.ctx, s.deviceFlowApp.ClientID, "127.0.0.1", "ticloud")
	require.Nil(s.T(), err)
	require.NotNil(s.T(), resp)

	invalidCases := []struct {
		caseName   string
		clientID   string
		deviceCode string
	}{
		{"empty", "", ""},
		{"invalid", "", "invalid"},
		{"not exist", "invalid", "not exist"},
		{"invalid device code", s.deviceFlowApp.ClientID, "invalid"},
		{"invalid client id", "invalid", resp.DeviceCode},
		{"authorization code grant type", s.authorizationCodeApp.ClientID, resp.DeviceCode},
	}
	for _, c := range invalidCases {
		s.T().Run(c.caseName, func(t *testing.T) {
			resp, err := s.oauth2Service.HandleDeviceTokenRequest(s.ctx, c.clientID, c.deviceCode)
			assert.NotNil(t, err)
			assert.Nil(t, resp)

			count, dbErr := entity.NewOauth2AccessTokenQuerySet(s.db).Count()
			assert.Equal(t, int64(0), count)
			assert.Nil(t, dbErr)
		})
	}
}

func (s *HandleDeviceTokenSuite) TestHandleDeviceTokenRequestExpired() {
	resp, err := s.oauth2Service.HandleDeviceAuthorizationRequest(s.ctx, s.deviceFlowApp.ClientID, "127.0.0.1", "ticloud")
	require.Nil(s.T(), err)
	require.NotNil(s.T(), resp)

	var grantObj entity.Oauth2DeviceGrant
	err = entity.NewOauth2DeviceGrantQuerySet(s.db).UserCodeEq(resp.UserCode).One(&grantObj)
	require.Nil(s.T(), err)

	// created_at is read only, so we need to update expires_in
	// TODO: update created_at manually, sleep is not a good way
	grantObj.ExpiresIn = 0
	err = grantObj.Update(s.db, []entity.Oauth2DeviceGrantDBSchemaField{
		entity.Oauth2DeviceGrantDBSchema.ExpiresIn,
	}...)
	require.Nil(s.T(), err)

	time.Sleep(time.Millisecond * 200)

	tokenResp, err := s.oauth2Service.HandleDeviceTokenRequest(s.ctx, s.deviceFlowApp.ClientID, resp.DeviceCode)
	assert.Equal(s.T(), errs.ErrOauth2ExpiredToken, err)
	assert.Nil(s.T(), tokenResp)

	count, dbErr := entity.NewOauth2AccessTokenQuerySet(s.db).Count()
	assert.Equal(s.T(), int64(0), count)
	assert.Nil(s.T(), dbErr)
}

func (s *HandleDeviceTokenSuite) TestHandleDeviceTokenRequestSlowDown() {
	resp, err := s.oauth2Service.HandleDeviceAuthorizationRequest(s.ctx, s.deviceFlowApp.ClientID, "127.0.0.1", "ticloud")
	require.Nil(s.T(), err)
	require.NotNil(s.T(), resp)

	var grantObj entity.Oauth2DeviceGrant
	err = entity.NewOauth2DeviceGrantQuerySet(s.db).UserCodeEq(resp.UserCode).One(&grantObj)
	require.Nil(s.T(), err)

	tokenResp, err := s.oauth2Service.HandleDeviceTokenRequest(s.ctx, s.deviceFlowApp.ClientID, resp.DeviceCode)
	assert.Equal(s.T(), errs.ErrOauth2AuthorizationPending, err)
	assert.Nil(s.T(), tokenResp)

	tokenResp, err = s.oauth2Service.HandleDeviceTokenRequest(s.ctx, s.deviceFlowApp.ClientID, resp.DeviceCode)
	assert.Equal(s.T(), errs.ErrOauth2SlowDown, err)
	assert.Nil(s.T(), tokenResp)

	count, dbErr := entity.NewOauth2AccessTokenQuerySet(s.db).Count()
	assert.Equal(s.T(), int64(0), count)
	assert.Nil(s.T(), dbErr)
}

func (s *HandleDeviceTokenSuite) TestHandleDeviceTokenRequestAccessDenied() {
	resp, err := s.oauth2Service.HandleDeviceAuthorizationRequest(s.ctx, s.deviceFlowApp.ClientID, "127.0.0.1", "ticloud")
	require.Nil(s.T(), err)
	require.NotNil(s.T(), resp)

	var grantObj entity.Oauth2DeviceGrant
	err = entity.NewOauth2DeviceGrantQuerySet(s.db).UserCodeEq(resp.UserCode).One(&grantObj)
	require.Nil(s.T(), err)

	grantObj.Status = entity.Oauth2GrantStatusDeny
	err = grantObj.Update(s.db, []entity.Oauth2DeviceGrantDBSchemaField{
		entity.Oauth2DeviceGrantDBSchema.Status,
	}...)
	require.Nil(s.T(), err)

	tokenResp, err := s.oauth2Service.HandleDeviceTokenRequest(s.ctx, s.deviceFlowApp.ClientID, resp.DeviceCode)
	assert.Equal(s.T(), errs.ErrOauth2AccessDenied, err)
	assert.Nil(s.T(), tokenResp)

	count, dbErr := entity.NewOauth2AccessTokenQuerySet(s.db).Count()
	assert.Equal(s.T(), int64(0), count)
	assert.Nil(s.T(), dbErr)
}

func (s *HandleDeviceTokenSuite) TestHandleDeviceTokenRequestRevoked() {
	resp, err := s.oauth2Service.HandleDeviceAuthorizationRequest(s.ctx, s.deviceFlowApp.ClientID, "127.0.0.1", "ticloud")
	require.Nil(s.T(), err)
	require.NotNil(s.T(), resp)

	var grantObj entity.Oauth2DeviceGrant
	err = entity.NewOauth2DeviceGrantQuerySet(s.db).UserCodeEq(resp.UserCode).One(&grantObj)
	require.Nil(s.T(), err)

	now := time.Now()
	grantObj.RevokedAt = &now
	err = grantObj.Update(s.db, []entity.Oauth2DeviceGrantDBSchemaField{
		entity.Oauth2DeviceGrantDBSchema.RevokedAt,
	}...)
	require.Nil(s.T(), err)

	tokenResp, err := s.oauth2Service.HandleDeviceTokenRequest(s.ctx, s.deviceFlowApp.ClientID, resp.DeviceCode)
	assert.Equal(s.T(), errs.ErrOauth2ExpiredToken, err)
	assert.Nil(s.T(), tokenResp)

	count, dbErr := entity.NewOauth2AccessTokenQuerySet(s.db).Count()
	assert.Equal(s.T(), int64(0), count)
	assert.Nil(s.T(), dbErr)
}

func (s *HandleDeviceTokenSuite) TestHandleDeviceTokenRequestSuccess() {
	resp, err := s.oauth2Service.HandleDeviceAuthorizationRequest(s.ctx, s.deviceFlowApp.ClientID, "127.0.0.1", "ticloud")
	require.Nil(s.T(), err)
	require.NotNil(s.T(), resp)

	var grantObj entity.Oauth2DeviceGrant
	err = entity.NewOauth2DeviceGrantQuerySet(s.db).UserCodeEq(resp.UserCode).One(&grantObj)
	require.Nil(s.T(), err)

	grantObj.Status = entity.Oauth2GrantStatusApprove
	grantObj.UserID = 1
	grantObj.UserEmail = "test@example.com"
	grantObj.OrgID = 1000
	grantObj.GrantAllProjects = true
	grantObj.ProjectIDs = entity.ProjectIDs{IDs: []uint64{}}
	grantObj.Scopes = "[]"
	err = grantObj.Update(s.db, []entity.Oauth2DeviceGrantDBSchemaField{
		entity.Oauth2DeviceGrantDBSchema.Status,
		entity.Oauth2DeviceGrantDBSchema.UserID,
		entity.Oauth2DeviceGrantDBSchema.UserEmail,
		entity.Oauth2DeviceGrantDBSchema.OrgID,
		entity.Oauth2DeviceGrantDBSchema.GrantAllProjects,
		entity.Oauth2DeviceGrantDBSchema.ProjectIDs,
		entity.Oauth2DeviceGrantDBSchema.Scopes,
	}...)
	require.Nil(s.T(), err)

	tokenResp, err := s.oauth2Service.HandleDeviceTokenRequest(s.ctx, s.deviceFlowApp.ClientID, resp.DeviceCode)
	assert.Nil(s.T(), err)
	require.NotNil(s.T(), tokenResp)
	assert.Equal(s.T(), tokenResp.TokenType, "Bearer")
	assert.Equal(s.T(), tokenResp.ExpiresIn, s.deviceFlowApp.AccessTokenExpiresIn)
	assert.NotEmpty(s.T(), tokenResp.AccessToken)
	assert.NotEmpty(s.T(), tokenResp.RefreshToken)

	encrptedAccessToken := crypto.NewPbdkf2Sha512Encryptor().Encrypt(
		tokenResp.AccessToken, "", 20000,
	)
	var akObj entity.Oauth2AccessToken
	dbErr := entity.NewOauth2AccessTokenQuerySet(s.db).AccessTokenEq(encrptedAccessToken).One(&akObj)
	assert.Nil(s.T(), dbErr)
	// should store the encrypted access token and refresh token in db
	assert.NotEqual(s.T(), tokenResp.AccessToken, akObj.AccessToken)
	assert.NotEqual(s.T(), tokenResp.RefreshToken, akObj.RefreshToken)
	assert.Equal(s.T(), tokenResp.ExpiresIn, akObj.AccessTokenExpiresIn)
	assert.Equal(s.T(), uint64(s.deviceFlowApp.ID), akObj.ApplicationID)
	assert.Equal(s.T(), uint64(1), akObj.UserID)
	assert.Equal(s.T(), "test@example.com", akObj.UserEmail)
	assert.Equal(s.T(), uint64(1000), akObj.OrgID)
	assert.Equal(s.T(), entity.Oauth2Scope{}, akObj.Scopes)
	assert.Equal(s.T(), entity.ProjectIDs{IDs: []uint64{}}, akObj.ProjectIDs)
	assert.Equal(s.T(), true, akObj.GrantAllProjects)
	assert.Equal(s.T(), entity.Oauth2DeviceCode, akObj.GrantType)
	assert.Nil(s.T(), akObj.RevokedAt)
}