package oidc

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/buildbuddy-io/buildbuddy/enterprise/server/testutil/enterprise_testenv"
	"github.com/buildbuddy-io/buildbuddy/server/tables"
	"github.com/buildbuddy-io/buildbuddy/server/util/authutil"
	"github.com/buildbuddy-io/buildbuddy/server/util/cookie"
	"github.com/buildbuddy-io/buildbuddy/server/util/status"
	"github.com/stretchr/testify/require"
	"golang.org/x/oauth2"
)

const (
	testIssuer        = "testIssuer"
	testSlug          = "test-slug"
	validJWT          = "validJWT"
	expiredJWT        = "expiredJWT"
	refreshedJWT      = "refreshedJWT"
	userID            = "1234"
	userName          = "Fluffy"
	userEmail         = "fluffy@cuteanimals.test"
	subID             = testIssuer + "/" + userID
	validRefreshToken = "abc123"
)

var (
	validUserToken = &userToken{
		issuer: testIssuer,
		Sub:    userID,
		Name:   userName,
	}
)

type fakeOidcAuthenticator struct {
}

func (f fakeOidcAuthenticator) getIssuer() string {
	return testIssuer
}

func (f fakeOidcAuthenticator) getSlug() string {
	return testSlug
}

func (f fakeOidcAuthenticator) authCodeURL(state string, opts ...oauth2.AuthCodeOption) (string, error) {
	return "https://auth.test", nil
}

func (f fakeOidcAuthenticator) exchange(ctx context.Context, code string, opts ...oauth2.AuthCodeOption) (*oauth2.Token, error) {
	return nil, status.UnimplementedError("not implemented")
}

func (f fakeOidcAuthenticator) verifyTokenAndExtractUser(ctx context.Context, jwt string, checkExpiry bool) (*userToken, error) {
	if jwt == validJWT || jwt == expiredJWT && !checkExpiry {
		return validUserToken, nil
	} else if jwt == expiredJWT {
		return nil, status.PermissionDeniedErrorf("expired JWT")
	}
	return nil, status.PermissionDeniedError("invalid JWT")
}

func (f fakeOidcAuthenticator) checkAccessToken(ctx context.Context, jwt, accessToken string) error {
	return nil
}

func (f fakeOidcAuthenticator) renewToken(ctx context.Context, refreshToken string) (*oauth2.Token, error) {
	if refreshToken == validRefreshToken {
		t := (&oauth2.Token{}).WithExtra(map[string]interface{}{"id_token": refreshedJWT})
		return t, nil
	}
	return nil, status.PermissionDeniedError("invalid refresh token")
}

func TestAuthenticateHTTPRequest(t *testing.T) {
	env := enterprise_testenv.GetCustomTestEnv(t, &enterprise_testenv.Options{})

	auth, err := newForTesting(context.Background(), env, &fakeOidcAuthenticator{})
	require.NoErrorf(t, err, "could not create authenticator")

	// JWT cookie not present, auth should fail.
	request, err := http.NewRequest(http.MethodGet, "/", strings.NewReader(""))
	require.NoErrorf(t, err, "could not create HTTP request")
	response := httptest.NewRecorder()
	authCtx := auth.AuthenticatedHTTPContext(response, request)
	requireAuthenticationError(t, authCtx)

	sessionID := "e34ff952-6ef0-4a35-ae3d-fe6166fe277e"
	// Valid JWT cookie, but user does not exist.
	request.AddCookie(&http.Cookie{Name: cookie.JWTCookie, Value: validJWT})
	request.AddCookie(&http.Cookie{Name: cookie.AuthIssuerCookie, Value: testIssuer})
	request.AddCookie(&http.Cookie{Name: cookie.SessionIDCookie, Value: sessionID})
	authCtx = auth.AuthenticatedHTTPContext(response, request)
	requireAuthenticationError(t, authCtx)

	// User information should still be populated (it's needed for user creation).
	require.Equal(t, validUserToken, authCtx.Value(contextUserKey), "context user details should match details returned by provider")

	// Create matching user, and token. Authentication should now succeed.
	user := &tables.User{
		UserID: userID,
		SubID:  subID,
		Email:  userEmail,
	}
	err = env.GetUserDB().InsertUser(context.Background(), user)
	require.NoError(t, err, "could not insert user")
	err = env.GetAuthDB().InsertOrUpdateUserSession(context.Background(), sessionID, &tables.Session{
		SessionID:    sessionID,
		SubID:        subID,
		AccessToken:  "access",
		RefreshToken: "refresh",
	})
	require.NoError(t, err, "could not insert token")
	authCtx = auth.AuthenticatedHTTPContext(response, request)
	requireAuthenticated(t, authCtx)
	require.Equal(t, validUserToken, authCtx.Value(contextUserKey), "context user details should match details returned by provider")

	// Send request with an expired JWT cookie.
	// DB doesn't have a refresh token so authentication should fail.
	request, err = http.NewRequest(http.MethodGet, "/", strings.NewReader(""))
	require.NoErrorf(t, err, "could not create HTTP request")
	request.AddCookie(&http.Cookie{Name: cookie.JWTCookie, Value: expiredJWT})
	request.AddCookie(&http.Cookie{Name: cookie.AuthIssuerCookie, Value: testIssuer})
	request.AddCookie(&http.Cookie{Name: cookie.SessionIDCookie, Value: sessionID})
	authCtx = auth.AuthenticatedHTTPContext(response, request)
	requireAuthenticationError(t, authCtx)

	// Insert a refresh token into the DB & auth should succeed.
	session := &tables.Session{RefreshToken: validRefreshToken}
	err = env.GetAuthDB().InsertOrUpdateUserSession(context.Background(), sessionID, session)
	require.NoError(t, err, "could not insert token")
	response = httptest.NewRecorder()
	authCtx = auth.AuthenticatedHTTPContext(response, request)
	requireAuthenticated(t, authCtx)
	newJwtCookie := getResponseCookie(response.Result(), cookie.JWTCookie)
	require.NotNil(t, newJwtCookie, "JWT cookie should be updated")
	require.Equal(t, refreshedJWT, newJwtCookie.Value, "JWT cookie should be updated")
	require.Equal(t, validUserToken, authCtx.Value(contextUserKey), "context user details should match details returned by provider")
}

func getResponseCookie(response *http.Response, name string) *http.Cookie {
	for _, c := range response.Cookies() {
		if c.Name == name {
			return c
		}
	}
	return nil
}

func requireAuthenticationError(t *testing.T, ctx context.Context) {
	err, _ := authutil.AuthErrorFromContext(ctx)
	require.NotNil(t, err, "context auth error key should be set")
	require.Nil(t, ctx.Value(authutil.ContextTokenStringKey), "context auth jwt token should not be set")
}

func requireAuthenticated(t *testing.T, ctx context.Context) {
	err, _ := authutil.AuthErrorFromContext(ctx)
	require.Nil(t, err, err)
	require.NotNil(t, ctx.Value(authutil.ContextTokenStringKey), "context auth jwt token should be set")
}
