package main

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"testing"

	pdktest "github.com/Kong/go-pdk/test"
	"github.com/stretchr/testify/require"
	baseapi "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/base"
	authv1 "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/core-server/auth/v1"
	userv1 "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/core-server/user/v1"
	userentityv1 "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/entity/core-server/user/v1"
	authconsts "github.com/tidbcloud/aisupport/packages/sdk-go/infra/auth/consts"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

func TestAccessTokenFromAuthHeader(t *testing.T) {
	t.Parallel()

	require.Equal(t, "token", accessTokenFromAuthHeader("Bearer token"))
	require.Equal(t, "token", accessTokenFromAuthHeader("bearer token"))
	require.Empty(t, accessTokenFromAuthHeader(""))
	require.Empty(t, accessTokenFromAuthHeader("Basic token"))
	require.Empty(t, accessTokenFromAuthHeader("Bearer"))
}

func TestAuthenticateRequest(t *testing.T) {
	t.Parallel()

	req := authenticateRequest(credential{accessToken: "access", apiToken: "api-token", apiKey: "api-key"})
	require.Equal(t, "access", req.GetAccessToken())
	require.Empty(t, req.GetApiToken())
	require.Empty(t, req.GetApiKey())

	req = authenticateRequest(credential{apiToken: "api-token", apiKey: "api-key"})
	require.Equal(t, "api-token", req.GetApiToken())
	require.Empty(t, req.GetApiKey())

	req = authenticateRequest(credential{apiKey: "api-key"})
	require.Equal(t, "api-key", req.GetApiKey())
}

func TestAuthorizeRequest(t *testing.T) {
	t.Parallel()

	req := authorizeRequest(&authv1.AuthenticateResponse{PrincipalProfile: &userv1.PrincipalProfile{
		OperatorType: authconsts.OperatorTypeUser,
		UserEmail:    "user@pingcap.com",
	}}, "/api/v1/test", "GET")
	require.Equal(t, "user@pingcap.com", req.GetUserEmail())
	require.Empty(t, req.GetApiKeyName())
	require.Equal(t, "/api/v1/test", req.GetUri())
	require.Equal(t, "GET", req.GetMethod())

	req = authorizeRequest(&authv1.AuthenticateResponse{PrincipalProfile: &userv1.PrincipalProfile{
		OperatorType: authconsts.OperatorTypeAPIKey,
		OperatorName: "key-name",
	}}, "/api/v1/test", "POST")
	require.Equal(t, "key-name", req.GetApiKeyName())
	require.Empty(t, req.GetUserEmail())

	req = authorizeRequest(&authv1.AuthenticateResponse{
		PrincipalProfile: &userv1.PrincipalProfile{
			OperatorType: authconsts.OperatorTypeAPIToken,
			UserEmail:    "owner@pingcap.com",
		},
		ApiToken: &userentityv1.APIToken{Privilege: userentityv1.APITokenPrivilege_API_TOKEN_PRIVILEGE_READONLY},
	}, "/api/v1/test", "GET")
	require.Equal(t, "owner@pingcap.com", req.GetUserEmail())
	require.Equal(t, userentityv1.APITokenPrivilege_API_TOKEN_PRIVILEGE_READONLY, req.GetApiTokenPrivilege())
}

func TestAccessWithBearerToken(t *testing.T) {
	t.Parallel()

	fake := &fakeAuthClient{
		authenticateResp: &authv1.AuthenticateResponse{PrincipalProfile: &userv1.PrincipalProfile{
			OperatorType: authconsts.OperatorTypeUser,
			OperatorName: "user@pingcap.com",
			UserEmail:    "user@pingcap.com",
			AggregatedRoleScope: &userv1.PrincipalRoleScope{
				AllowedTenantIds: []uint64{1001},
			},
			Roles: map[string]*userv1.PrincipalRoleScope{
				"role-user": {AllowedTenantIds: []uint64{1001}},
			},
		}},
		authorizeResp: &authv1.AuthorizeResponse{
			Authorized: true,
		},
	}
	env := runAccess(t, fake, pdktest.Request{
		Method: "GET",
		Url:    "https://example.com/api/v1/test?tenant_id=1001",
		Headers: http.Header{
			authconsts.HeaderAuthorization: {"Bearer access-token"},
		},
	})

	require.Equal(t, "access-token", fake.authenticateRequests[0].GetAccessToken())
	require.Equal(t, "/api/v1/test", fake.authorizeRequests[0].GetUri())
	require.Equal(t, "GET", fake.authorizeRequests[0].GetMethod())
	require.Equal(t, "user@pingcap.com", fake.authorizeRequests[0].GetUserEmail())
	require.Equal(t, authconsts.OperatorTypeUser, env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderOperatorType)))
	require.Equal(t, "user@pingcap.com", env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderOperator)))
	require.Equal(t, "user@pingcap.com", env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderUserEmail)))
	require.Equal(t, "10.10.10.1", fake.authenticateMetadata[0].Get(strings.ToLower(authconsts.HeaderClientIP))[0])
	require.Equal(t, "10.10.10.1", env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderClientIP)))
	require.JSONEq(t, `{
		"operator_type":"user",
		"operator_name":"user@pingcap.com",
		"user_email":"user@pingcap.com",
		"aggregated_role_scope":{"allowed_tenant_ids":[1001]},
		"roles":{"role-user":{"allowed_tenant_ids":[1001]}}
	}`, env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderPrincipalProfile)))
	require.Equal(t, authconsts.OperatorTypeUser, env.ServiceReq.Headers.Get(authconsts.HeaderOperatorType))
	require.Equal(t, "user@pingcap.com", env.ServiceReq.Headers.Get(authconsts.HeaderOperator))
	require.Equal(t, "user@pingcap.com", env.ServiceReq.Headers.Get(authconsts.HeaderUserEmail))
	require.JSONEq(t, `{
		"operator_type":"user",
		"operator_name":"user@pingcap.com",
		"user_email":"user@pingcap.com",
		"aggregated_role_scope":{"allowed_tenant_ids":[1001]},
		"roles":{"role-user":{"allowed_tenant_ids":[1001]}}
	}`, env.ServiceReq.Headers.Get(authconsts.HeaderPrincipalProfile))
	require.Empty(t, env.ServiceReq.Headers.Get("Grpc-Metadata-X-Role-Scope"))
}

func TestAccessWithAPIToken(t *testing.T) {
	t.Parallel()

	fake := &fakeAuthClient{
		authenticateResp: &authv1.AuthenticateResponse{
			ApiToken: &userentityv1.APIToken{Privilege: userentityv1.APITokenPrivilege_API_TOKEN_PRIVILEGE_READONLY},
			PrincipalProfile: &userv1.PrincipalProfile{
				OperatorType:        authconsts.OperatorTypeAPIToken,
				OperatorName:        "owner@pingcap.com",
				UserEmail:           "owner@pingcap.com",
				AggregatedRoleScope: &userv1.PrincipalRoleScope{},
				Roles:               map[string]*userv1.PrincipalRoleScope{},
			},
		},
		authorizeResp: &authv1.AuthorizeResponse{Authorized: true},
	}
	env := runAccess(t, fake, pdktest.Request{
		Method: "GET",
		Url:    "https://example.com/api/v1/test",
		Headers: http.Header{
			authconsts.HeaderAPIToken: {"api-token"},
		},
	})

	require.Equal(t, "api-token", fake.authenticateRequests[0].GetApiToken())
	require.Equal(t, "owner@pingcap.com", fake.authorizeRequests[0].GetUserEmail())
	require.Equal(t, userentityv1.APITokenPrivilege_API_TOKEN_PRIVILEGE_READONLY, fake.authorizeRequests[0].GetApiTokenPrivilege())
	require.Equal(t, authconsts.OperatorTypeAPIToken, env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderOperatorType)))
	require.Equal(t, "owner@pingcap.com", env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderOperator)))
	require.Equal(t, "owner@pingcap.com", env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderUserEmail)))
	require.Equal(t, "10.10.10.1", fake.authenticateMetadata[0].Get(strings.ToLower(authconsts.HeaderClientIP))[0])
	require.Equal(t, "10.10.10.1", env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderClientIP)))
	require.JSONEq(t, `{
		"operator_type":"api_token",
		"operator_name":"owner@pingcap.com",
		"user_email":"owner@pingcap.com",
		"aggregated_role_scope":{}
	}`, env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderPrincipalProfile)))
}

func TestAccessWithAPIKey(t *testing.T) {
	t.Parallel()

	fake := &fakeAuthClient{
		authenticateResp: &authv1.AuthenticateResponse{
			PrincipalProfile: &userv1.PrincipalProfile{
				OperatorType: authconsts.OperatorTypeAPIKey,
				OperatorName: "key-name",
				UserEmail:    "owner@pingcap.com",
				AggregatedRoleScope: &userv1.PrincipalRoleScope{
					AllowedTenantIds: []uint64{2001},
				},
				Roles: map[string]*userv1.PrincipalRoleScope{
					"role-key": {AllowedTenantIds: []uint64{2001}},
				},
			},
		},
		authorizeResp: &authv1.AuthorizeResponse{
			Authorized: true,
		},
	}
	env := runAccess(t, fake, pdktest.Request{
		Method: "GET",
		Url:    "https://example.com/api/v1/test",
		Headers: http.Header{
			authconsts.HeaderAPIKey: {"api-key"},
		},
	})

	require.Equal(t, "api-key", fake.authenticateRequests[0].GetApiKey())
	require.Equal(t, "key-name", fake.authorizeRequests[0].GetApiKeyName())
	require.Equal(t, authconsts.OperatorTypeAPIKey, env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderOperatorType)))
	require.Equal(t, "key-name", env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderOperator)))
	require.Equal(t, "owner@pingcap.com", env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderUserEmail)))
	require.Equal(t, "10.10.10.1", fake.authenticateMetadata[0].Get(strings.ToLower(authconsts.HeaderClientIP))[0])
	require.Equal(t, "10.10.10.1", env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderClientIP)))
	require.Empty(t, env.ServiceReq.Headers.Get("X-API-Key-Name"))
	require.Empty(t, env.ServiceReq.Headers.Get("X-API-Key-Owner"))
	require.JSONEq(t, `{
		"operator_type":"api_key",
		"operator_name":"key-name",
		"user_email":"owner@pingcap.com",
		"aggregated_role_scope":{"allowed_tenant_ids":[2001]},
		"roles":{"role-key":{"allowed_tenant_ids":[2001]}}
	}`, env.ServiceReq.Headers.Get(authconsts.GRPCMetadataKeyFromHTTPHeader(authconsts.HeaderPrincipalProfile)))
}

func TestAccessCredentialPrecedence(t *testing.T) {
	t.Parallel()

	fake := &fakeAuthClient{
		authenticateResp: &authv1.AuthenticateResponse{PrincipalProfile: &userv1.PrincipalProfile{
			OperatorType: authconsts.OperatorTypeUser,
			OperatorName: "user@pingcap.com",
			UserEmail:    "user@pingcap.com",
		}},
		authorizeResp: &authv1.AuthorizeResponse{Authorized: true},
	}
	runAccess(t, fake, pdktest.Request{
		Method: "GET",
		Url:    "https://example.com/api/v1/test",
		Headers: http.Header{
			authconsts.HeaderAuthorization: {"Bearer access-token"},
			authconsts.HeaderAPIToken:      {"api-token"},
			authconsts.HeaderAPIKey:        {"api-key"},
		},
	})

	require.Equal(t, "access-token", fake.authenticateRequests[0].GetAccessToken())
	require.Empty(t, fake.authenticateRequests[0].GetApiToken())
	require.Empty(t, fake.authenticateRequests[0].GetApiKey())
}

func TestAccessMissingCredential(t *testing.T) {
	t.Parallel()

	env := runAccess(t, &fakeAuthClient{}, pdktest.Request{
		Method: "GET",
		Url:    "https://example.com/api/v1/test",
		Headers: http.Header{
			"Origin": {"https://ticket-manager.example.com"},
		},
	})

	require.Equal(t, http.StatusUnauthorized, env.ClientRes.Status)
	require.Equal(t, "application/json", env.ClientRes.Headers.Get("Content-Type"))
	require.Equal(t, "https://ticket-manager.example.com", env.ClientRes.Headers.Get("Access-Control-Allow-Origin"))
	require.Equal(t, "true", env.ClientRes.Headers.Get("Access-Control-Allow-Credentials"))
	require.Equal(t, "Origin", env.ClientRes.Headers.Get("Vary"))
	resp := decodeBaseResponse(t, env)
	require.Equal(t, http.StatusUnauthorized, resp.ErrCode)
	require.Equal(t, "miss credential", resp.ErrMessage)
}

func TestAccessAuthenticateError(t *testing.T) {
	t.Parallel()

	env := runAccess(t, &fakeAuthClient{authenticateErr: errors.New("authn failed")}, pdktest.Request{
		Method: "GET",
		Url:    "https://example.com/api/v1/test",
		Headers: http.Header{
			authconsts.HeaderAuthorization: {"Bearer access-token"},
		},
	})

	require.Equal(t, http.StatusUnauthorized, env.ClientRes.Status)
	resp := decodeBaseResponse(t, env)
	require.Equal(t, http.StatusUnauthorized, resp.ErrCode)
	require.Equal(t, "authn failed", resp.ErrMessage)
}

func TestAccessAuthenticatePermissionDenied(t *testing.T) {
	t.Parallel()

	env := runAccess(t, &fakeAuthClient{authenticateErr: status.Error(codes.PermissionDenied, "ip not allowed")}, pdktest.Request{
		Method: "GET",
		Url:    "https://example.com/api/v1/test",
		Headers: http.Header{
			authconsts.HeaderAPIToken: {"api-token"},
		},
	})

	require.Equal(t, http.StatusForbidden, env.ClientRes.Status)
	resp := decodeBaseResponse(t, env)
	require.Equal(t, http.StatusForbidden, resp.ErrCode)
	require.Equal(t, "rpc error: code = PermissionDenied desc = ip not allowed", resp.ErrMessage)
}

func TestAccessAuthorizeDenied(t *testing.T) {
	t.Parallel()

	env := runAccess(t, &fakeAuthClient{
		authenticateResp: &authv1.AuthenticateResponse{PrincipalProfile: &userv1.PrincipalProfile{
			OperatorType: authconsts.OperatorTypeUser,
			OperatorName: "user@pingcap.com",
			UserEmail:    "user@pingcap.com",
		}},
		authorizeResp: &authv1.AuthorizeResponse{Authorized: false},
	}, pdktest.Request{
		Method: "GET",
		Url:    "https://example.com/api/v1/test",
		Headers: http.Header{
			authconsts.HeaderAuthorization: {"Bearer access-token"},
		},
	})

	require.Equal(t, http.StatusForbidden, env.ClientRes.Status)
	resp := decodeBaseResponse(t, env)
	require.Equal(t, http.StatusForbidden, resp.ErrCode)
	require.Equal(t, "permission denied", resp.ErrMessage)
}

func runAccess(t *testing.T, client *fakeAuthClient, req pdktest.Request) *pdktest.TestEnv {
	t.Helper()

	env, err := pdktest.New(t, req)
	require.NoError(t, err)

	env.DoAccess(&Config{
		RequestTimeoutMS: 1000,
		authClient:       client,
	})
	return env
}

func decodeBaseResponse(t *testing.T, env *pdktest.TestEnv) baseapi.BaseResponse {
	t.Helper()

	var resp baseapi.BaseResponse
	require.NoError(t, json.Unmarshal(env.ClientRes.Body, &resp))
	return resp
}

type fakeAuthClient struct {
	authenticateRequests []*authv1.AuthenticateRequest
	authenticateMetadata []metadata.MD
	authorizeRequests    []*authv1.AuthorizeRequest

	authenticateResp *authv1.AuthenticateResponse
	authorizeResp    *authv1.AuthorizeResponse
	authenticateErr  error
	authorizeErr     error
}

func (f *fakeAuthClient) Authenticate(ctx context.Context, req *authv1.AuthenticateRequest, _ ...grpc.CallOption) (*authv1.AuthenticateResponse, error) {
	f.authenticateRequests = append(f.authenticateRequests, req)
	md, _ := metadata.FromOutgoingContext(ctx)
	f.authenticateMetadata = append(f.authenticateMetadata, md)
	return f.authenticateResp, f.authenticateErr
}

func (f *fakeAuthClient) Authorize(_ context.Context, req *authv1.AuthorizeRequest, _ ...grpc.CallOption) (*authv1.AuthorizeResponse, error) {
	f.authorizeRequests = append(f.authorizeRequests, req)
	return f.authorizeResp, f.authorizeErr
}