package oauth2_test

import (
	"errors"
	"testing"
	"time"

	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"

	"github.com/tidbcloud/account/idl/pbgen/proto/account"
	"github.com/tidbcloud/auth-server/commons/crypto"
	"github.com/tidbcloud/auth-server/commons/errs"
	"github.com/tidbcloud/auth-server/internal/entity"
)

type AuthorizationCodeSuite struct {
	Oauth2Suite
}

func TestAuthorizationCodeSuite(t *testing.T) {
	suite.Run(t, new(AuthorizationCodeSuite))
}

func (s *AuthorizationCodeSuite) TestAuthorizationCodeValidateInvalidRequest() {
	s.T().Run("app not found", func(t *testing.T) {
		req, err := s.oauth2Service.InnerAuthorizationCodeRequestValidate(s.ctx, s.testUser, "invalid", "code", s.authorizationCodeRawRedirectURI, "state", []string{"org:member", "project:owner"})
		assert.NotNil(t, err)
		assert.Equal(t, errs.ErrOauth2InvalidClient.Error(), err.Error())
		assert.Nil(t, req)
	})

	s.T().Run("app not authorization code", func(t *testing.T) {
		req, err := s.oauth2Service.InnerAuthorizationCodeRequestValidate(s.ctx, s.testUser, s.deviceFlowApp.ClientID, "code", s.authorizationCodeRawRedirectURI, "state", []string{"org:member", "project:owner"})
		assert.NotNil(t, err)
		assert.Equal(t, errs.ErrOauth2InvalidClient.Error(), err.Error())
		assert.Nil(t, req)
	})

	s.T().Run("invalid response type", func(t *testing.T) {
		req, err := s.oauth2Service.InnerAuthorizationCodeRequestValidate(s.ctx, s.testUser, s.authorizationCodeApp.ClientID, "invalid", s.authorizationCodeRawRedirectURI, "state", []string{"org:member", "project:owner"})
		assert.NotNil(t, err)
		assert.Equal(t, errs.ErrOauth2UnsupportedResponseType.Error(), err.Error())
		assert.NotNil(t, req)
		assert.NotNil(t, req.RedirectURI)
		assert.Equal(t, s.authorizationCodeApp.ClientID, req.ClientID)
	})

	s.T().Run("invalid redirect uri", func(t *testing.T) {
		req, err := s.oauth2Service.InnerAuthorizationCodeRequestValidate(s.ctx, s.testUser, s.authorizationCodeApp.ClientID, "code", "https://invalid.com", "state", []string{"org:member", "project:owner"})
		assert.NotNil(t, err)
		assert.Equal(t, errs.ErrOauth2RedirectURIMismatch.Error(), err.Error())
		assert.NotNil(t, req)
		assert.Equal(t, s.authorizationCodeApp.ClientID, req.ClientID)
	})

	s.T().Run("ivalid redirect uri with subpath", func(t *testing.T) {
		req, err := s.oauth2Service.InnerAuthorizationCodeRequestValidate(s.ctx, s.testUser, s.authorizationCodeApp.ClientID, "code", s.authorizationCodeRawRedirectURI+"/subpath", "state", []string{"org:member", "project:owner"})
		assert.NotNil(t, err)
		assert.Equal(t, errs.ErrOauth2RedirectURIMismatch.Error(), err.Error())
		assert.NotNil(t, req)
		assert.Nil(t, req.RedirectURI)
		assert.Equal(t, s.authorizationCodeApp.ClientID, req.ClientID)
	})

	s.T().Run("invalid scope 1", func(t *testing.T) {
		req, err := s.oauth2Service.InnerAuthorizationCodeRequestValidate(s.ctx, s.testUser, s.authorizationCodeApp.ClientID, "code", s.authorizationCodeRawRedirectURI, "state", []string{"org:invalid", "project:invalid"})
		assert.NotNil(t, err)
		assert.Equal(t, errs.ErrOauth2InvalidScope.Error(), err.Error())
		assert.NotNil(t, req)
		assert.Equal(t, s.authorizationCodeApp.ClientID, req.ClientID)
	})

	s.T().Run("invalid scope 2", func(t *testing.T) {
		req, err := s.oauth2Service.InnerAuthorizationCodeRequestValidate(s.ctx, s.testUser, s.authorizationCodeApp.ClientID, "code", s.authorizationCodeRawRedirectURI, "state", []string{"org:owner", "project:owner"})
		assert.NotNil(t, err)
		assert.Equal(t, errs.ErrOauth2InvalidScope.Error(), err.Error())
		assert.NotNil(t, req)
		assert.NotNil(t, req.RedirectURI)
		assert.Equal(t, s.authorizationCodeApp.ClientID, req.ClientID)
	})
}

func (s *AuthorizationCodeSuite) TestAuthorizationCodeValidateValidRequest() {
	s.T().Run("valid request with same redirect uri", func(t *testing.T) {
		req, err := s.oauth2Service.InnerAuthorizationCodeRequestValidate(s.ctx, s.testUser, s.authorizationCodeApp.ClientID, "code", s.authorizationCodeRawRedirectURI, "state", []string{"org:member", "project:owner"})
		assert.Nil(t, err)
		assert.NotNil(t, req)
		assert.NotNil(t, req.RedirectURI)
		assert.Equal(t, s.authorizationCodeApp.ClientID, req.ClientID)
		assert.Equal(t, s.authorizationCodeRawRedirectURI, req.RedirectURI.String())
	})

	s.T().Run("valid request with empty redirect uri", func(t *testing.T) {
		req, err := s.oauth2Service.InnerAuthorizationCodeRequestValidate(s.ctx, s.testUser, s.authorizationCodeApp.ClientID, "code", "", "state", []string{"org:member", "project:owner"})
		assert.Nil(t, err)
		assert.NotNil(t, req)
		assert.NotNil(t, req.RedirectURI)
		assert.Equal(t, s.authorizationCodeApp.ClientID, req.ClientID)
		assert.Equal(t, s.authorizationCodeRawRedirectURI, req.RedirectURI.String())
	})

	s.T().Run("valid request with empty scope", func(t *testing.T) {
		req, err := s.oauth2Service.InnerAuthorizationCodeRequestValidate(s.ctx, s.testUser, s.authorizationCodeApp.ClientID, "code", s.authorizationCodeRawRedirectURI, "state", []string{})
		assert.Nil(t, err)
		assert.NotNil(t, req)
		assert.NotNil(t, req.RedirectURI)
		assert.Equal(t, s.authorizationCodeApp.ClientID, req.ClientID)
		assert.Equal(t, s.authorizationCodeRawRedirectURI, req.RedirectURI.String())
	})

	s.T().Run("valid request with minor project scope", func(t *testing.T) {
		req, err := s.oauth2Service.InnerAuthorizationCodeRequestValidate(s.ctx, s.testUser, s.authorizationCodeApp.ClientID, "code", s.authorizationCodeRawRedirectURI, "state", []string{"org:member", "project:dev"})
		assert.Nil(t, err)
		assert.NotNil(t, req)
		assert.NotNil(t, req.RedirectURI)
		assert.Equal(t, s.authorizationCodeApp.ClientID, req.ClientID)
		assert.Equal(t, s.authorizationCodeRawRedirectURI, req.RedirectURI.String())
	})
}

func (s *AuthorizationCodeSuite) TestAuthorizationCodeApproveNotInOrg() {
	s.accountMockProvider.EXPECT().VerifyUserOrgAndProjects(gomock.Any(), &account.VerifyUserOrgAndProjectsReq{
		UserId:      s.testUser.Id,
		OrgId:       1000,
		ProjectIds:  []uint64{1, 2},
		OrgRole:     s.authorizationCodeReq.OrgRole,
		ProjectRole: s.authorizationCodeReq.ProjectRole,
	}).Return(
		&account.VerifyUserOrgAndProjectsResp{Result: false}, nil,
	)

	grant, redirectURL, err := s.oauth2Service.InnerAuthorizationCodeApprove(s.ctx, s.testUser, s.authorizationCodeReq, 1000, []uint64{1, 2})
	assert.NotNil(s.T(), err)
	assert.Equal(s.T(), "user has no permission to access org or projects", err.Error())
	assert.Nil(s.T(), grant)
	assert.Nil(s.T(), redirectURL)
}

func (s *AuthorizationCodeSuite) TestAuthorizationCodeApproveSuccess() {
	s.accountMockProvider.EXPECT().VerifyUserOrgAndProjects(gomock.Any(), &account.VerifyUserOrgAndProjectsReq{
		UserId:      s.testUser.Id,
		OrgId:       1000,
		ProjectIds:  []uint64{1, 2},
		OrgRole:     s.authorizationCodeReq.OrgRole,
		ProjectRole: s.authorizationCodeReq.ProjectRole,
	}).Return(
		&account.VerifyUserOrgAndProjectsResp{Result: true}, nil,
	)

	grant, redirectURL, err := s.oauth2Service.InnerAuthorizationCodeApprove(s.ctx, s.testUser, s.authorizationCodeReq, 1000, []uint64{1, 2})
	assert.Nil(s.T(), err)
	assert.NotNil(s.T(), grant)
	assert.NotNil(s.T(), redirectURL)
	assert.Equal(s.T(), s.testUser.Id, grant.UserID)
	assert.Equal(s.T(), s.testUser.Email, grant.UserEmail)
	assert.Equal(s.T(), uint64(1000), grant.OrgID)
	assert.Equal(s.T(), entity.ProjectIDs{IDs: []uint64{1, 2}}, grant.ProjectIDs)
	assert.Equal(s.T(), entity.Oauth2Scope{Org: "org:member", Project: "project:owner"}, grant.Scopes)
	assert.Equal(s.T(), s.authorizationCodeApp.ID, entity.ID(grant.ApplicationID))
	assert.Equal(s.T(), 30, len(grant.Code))
	assert.Equal(s.T(), redirectURL.String(), s.authorizationCodeRawRedirectURI)
}

func (s *AuthorizationCodeSuite) TestAuthorizationCodeApproveFailed() {
	s.accountMockProvider.EXPECT().VerifyUserOrgAndProjects(gomock.Any(), &account.VerifyUserOrgAndProjectsReq{
		UserId:      s.testUser.Id,
		OrgId:       1000,
		ProjectIds:  []uint64{1, 2},
		OrgRole:     s.authorizationCodeReq.OrgRole,
		ProjectRole: s.authorizationCodeReq.ProjectRole,
	}).Return(
		&account.VerifyUserOrgAndProjectsResp{Result: false}, errors.New("some error"),
	)
	grant, redirectURL, err := s.oauth2Service.InnerAuthorizationCodeApprove(s.ctx, s.testUser, s.authorizationCodeReq, 1000, []uint64{1, 2})
	assert.NotNil(s.T(), err)
	assert.Nil(s.T(), grant)
	assert.Nil(s.T(), redirectURL)
}

func (s *AuthorizationCodeSuite) TestCreateAccessTokenByCodeGrantSyncToAccountFailed() {
	grant := s.createCodeGrant()
	s.accountMockProvider.EXPECT().SaveAuthorizationData(gomock.Any(), gomock.Any()).Return(
		&account.SaveAuthorizationDataResp{}, errors.New("new error"),
	).Times(1)
	resp, err := s.oauth2Service.CreateAccessTokenByCodeGrant(
		s.ctx, s.authorizationCodeApp.ClientID, s.authorizationCodeAppSecret, grant.Code, s.authorizationCodeRawRedirectURI,
	)
	assert.NotNil(s.T(), err)
	assert.Nil(s.T(), resp)

	// should not delete the code grant
	var grantObj entity.Oauth2CodeGrant
	dbErr := entity.NewOauth2CodeGrantQuerySet(s.db).CodeEq(grant.Code).One(&grantObj)
	assert.Nil(s.T(), dbErr)
}

func (s *AuthorizationCodeSuite) TestCreateAccessTokenByCodeGrantSuccess() {
	grant := s.createCodeGrant()
	s.accountMockProvider.EXPECT().SaveAuthorizationData(gomock.Any(), gomock.Any()).Return(&account.SaveAuthorizationDataResp{}, nil).Times(1)
	resp, err := s.oauth2Service.CreateAccessTokenByCodeGrant(
		s.ctx, s.authorizationCodeApp.ClientID, s.authorizationCodeAppSecret, grant.Code, s.authorizationCodeRawRedirectURI,
	)
	assert.Nil(s.T(), err)
	assert.NotNil(s.T(), resp)

	encrptedAccessToken := crypto.NewPbdkf2Sha512Encryptor().Encrypt(
		resp.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(), resp.AccessToken, akObj.AccessToken)
	assert.NotEqual(s.T(), resp.RefreshToken, akObj.RefreshToken)
	assert.Equal(s.T(), resp.ExpiresIn, akObj.AccessTokenExpiresIn)
	assert.Equal(s.T(), uint64(s.authorizationCodeApp.ID), akObj.ApplicationID)
	assert.Equal(s.T(), s.testUser.Id, akObj.UserID)
	assert.Equal(s.T(), s.testUser.Email, akObj.UserEmail)
	assert.Equal(s.T(), uint64(1000), akObj.OrgID)
	assert.Equal(s.T(), entity.Oauth2Scope{Org: "org:member", Project: "project:owner"}, akObj.Scopes)
	assert.Equal(s.T(), entity.ProjectIDs{IDs: []uint64{1, 2}}, akObj.ProjectIDs)
	assert.Equal(s.T(), false, akObj.GrantAllProjects)
	assert.Equal(s.T(), entity.Oauth2AuthorizationCode, akObj.GrantType)
	assert.Nil(s.T(), akObj.RevokedAt)

	// should delete the code grant
	var grantObj entity.Oauth2CodeGrant
	dbErr = entity.NewOauth2CodeGrantQuerySet(s.db).CodeEq(grant.Code).One(&grantObj)
	assert.NotNil(s.T(), dbErr)
}

func (s *AuthorizationCodeSuite) TestCreateAccessTokenByCodeGrantInvalidRequest() {
	grant := s.createCodeGrant()
	s.T().Run("invalid client id", func(t *testing.T) {
		resp, err := s.oauth2Service.CreateAccessTokenByCodeGrant(s.ctx, "invalid", s.authorizationCodeAppSecret, grant.Code, s.authorizationCodeRawRedirectURI)
		assert.NotNil(t, err)
		assert.Nil(t, resp)
	})

	s.T().Run("invalid client secret", func(t *testing.T) {
		resp, err := s.oauth2Service.CreateAccessTokenByCodeGrant(s.ctx, s.authorizationCodeApp.ClientID, "invalid", grant.Code, s.authorizationCodeRawRedirectURI)
		assert.NotNil(t, err)
		assert.Nil(t, resp)
	})

	s.T().Run("invalid code", func(t *testing.T) {
		resp, err := s.oauth2Service.CreateAccessTokenByCodeGrant(s.ctx, s.authorizationCodeApp.ClientID, s.authorizationCodeAppSecret, "invalid", s.authorizationCodeRawRedirectURI)
		assert.NotNil(t, err)
		assert.Nil(t, resp)
	})

	s.T().Run("mismatch redirect uri", func(t *testing.T) {
		resp, err := s.oauth2Service.CreateAccessTokenByCodeGrant(s.ctx, s.authorizationCodeApp.ClientID, s.authorizationCodeAppSecret, grant.Code, "invalid")
		assert.NotNil(t, err)
		assert.Nil(t, resp)
	})
}

func (s *AuthorizationCodeSuite) TestCreateAccessTokenByCodeGrantExpired() {
	grant := s.createCodeGrant()
	grant.ExpiresIn = 0
	err := grant.Update(s.db, []entity.Oauth2CodeGrantDBSchemaField{
		entity.Oauth2CodeGrantDBSchema.ExpiresIn,
	}...)
	require.Nil(s.T(), err)
	time.Sleep(time.Millisecond * 200)
	resp, err := s.oauth2Service.CreateAccessTokenByCodeGrant(s.ctx, s.authorizationCodeApp.ClientID, s.authorizationCodeAppSecret, grant.Code, s.authorizationCodeRawRedirectURI)
	assert.NotNil(s.T(), err)
	assert.Nil(s.T(), resp)
	assert.Equal(s.T(), errs.ErrOauth2InvalidGrant.Error(), err.Error())
}