package debug

import (
	"context"
	"net/http"
	"time"

	"github.com/form3tech-oss/jwt-go"
	"github.com/pingcap/errcode"
	"github.com/pingcap/errors"
	grpc_http "github.com/tidbcloud/metrics-server/central/grpc/http"
	pkg_http "github.com/tidbcloud/metrics-server/pkg/grpc/http"
	"github.com/tidbcloud/pkgv2/auth/util"
)

const (
	sessionTimeout                    = time.Hour * 48
	sessionCookieName                 = "debug.session_id"
	authByID                          = "debugSession"
	ctxKeyUser        util.ContextKey = "DebugSessionToken"
)

// The debug session is currently implemented by passing a JWT token in the HTTP header. The JWT token contains
// basic user information. The session is not managed in the database.

type Session struct {
	Email     string `json:"email"`
	Username  string `json:"username"`
	CSRFToken string `json:"csrfToken"`
}

type sessionTokenBody struct {
	jwt.StandardClaims
	Session *Session `json:"user"`
}

type SessionManager struct {
	// TODO: This session manager can be improved to be a general session manager,
	// not restricted to debug portal sessions.
	sessionSecret []byte
}

// NewSession creates a new session and returns the session id.
func (sm *SessionManager) NewSession(session *Session) (string, error) {
	// Currently the session key is simply generated by using a JWT token, storing session info in plain text.
	// In future it can be changed to store content in the database.
	tokenBody := sessionTokenBody{
		Session: session,
		StandardClaims: jwt.StandardClaims{
			ExpiresAt: time.Now().Add(sessionTimeout).Unix(),
		},
	}
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, tokenBody)
	tokenString, err := token.SignedString(sm.sessionSecret)
	if err != nil {
		return "", errcode.NewInternalErr(err)
	}
	return tokenString, nil
}

// GetSession gets the session user from the session id.
func (sm *SessionManager) getSession(sessionID string) (*Session, error) {
	if sessionID == "" {
		return nil, errcode.NewNotAuthenticatedErr(errors.New("sign in is required"))
	}
	token, err := jwt.ParseWithClaims(
		sessionID,
		&sessionTokenBody{},
		func(token *jwt.Token) (interface{}, error) {
			return sm.sessionSecret, nil
		})
	if err != nil {
		return nil, errcode.NewNotAuthenticatedErr(errors.Annotate(err, "invalid session key"))
	}
	tokenBody, ok := token.Claims.(*sessionTokenBody)
	if !ok || !token.Valid {
		return nil, errcode.NewNotAuthenticatedErr(errors.Annotate(err, "session is expired"))
	}
	return tokenBody.Session, nil
}

// MWTryAuth tries to authenticate the HTTP request according to the cookie.
// As a middleware, after a success auth, `AuthProcessedByKey` will be set.
// No error will be thrown when cookie does not exist, to allow other authentication methods to work.
// See `MWRequireAuth` for auth failure handling.
func (sm *SessionManager) MWTryAuth(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Don't authenticate an OPTIONS request
		if r.Method == http.MethodOptions {
			next.ServeHTTP(w, r)
			return
		}

		if r.Context().Value(util.AuthProcessedByKey) != nil {
			// Processed by other auth middlewares, skip.
			next.ServeHTTP(w, r)
			return
		}

		cookie, err := r.Cookie(sessionCookieName)
		if err != nil || cookie.Value == "" {
			next.ServeHTTP(w, r)
			return
		}

		sessionUser, err := sm.getSession(cookie.Value)
		if err != nil {
			authErr := errcode.NewNotAuthenticatedErr(err)
			code := errcode.CodeChain(authErr)
			http.Error(w, authErr.Error(), code.Code().HTTPCode())
			return
		}

		newCtx := context.WithValue(r.Context(), util.AuthProcessedByKey, authByID)
		newCtx = wrapContext(newCtx, sessionUser)
		newRequest := r.WithContext(newCtx)
		*r = *newRequest

		next.ServeHTTP(w, r)
	})
}

// GrpcGetDebugSessionUser authenticates the gRPC request according to the cookie.
// It also checks CSRF header `x-csrf-token`.
// Auth failures will result in errors.
// Currently it only checks whether session is valid. It does not support RBAC.
func (sm *SessionManager) GrpcGetDebugSessionUser(grpcCtx context.Context) (*Session, error) {
	sessionID := grpc_http.Cookie(grpcCtx, sessionCookieName)
	s, err := sm.getSession(sessionID)
	if err != nil {
		return nil, err
	}

	// CSRF token only need to be checked for non GET requests.
	method := grpc_http.HttpMethod(grpcCtx)
	if method != http.MethodGet {
		csrfToken := pkg_http.FirstHeader(grpcCtx, "x-csrf-token")
		if csrfToken != s.CSRFToken {
			return nil, errcode.NewForbiddenErr(errors.New("missing or incorrect x-csrf-token"))
		}
	}
	return s, nil
}

func GrpcSetCookie(grpcCtx context.Context, sessionKey string) error {
	cookie := http.Cookie{
		HttpOnly: true,
		Secure:   true,
		Name:     sessionCookieName,
		Value:    sessionKey,
		Path:     "/",
		MaxAge:   int(sessionTimeout / time.Second),
		SameSite: http.SameSiteNoneMode,
	}
	return grpc_http.SetHeader(grpcCtx, "Set-Cookie", cookie.String())
}

func GrpcClearCookie(grpcCtx context.Context) error {
	cookie := http.Cookie{
		HttpOnly: true,
		Secure:   true,
		Name:     sessionCookieName,
		Value:    "",
		Path:     "/",
		MaxAge:   -1,
		SameSite: http.SameSiteNoneMode,
	}
	return grpc_http.SetHeader(grpcCtx, "Set-Cookie", cookie.String())
}

func wrapContext(ctx context.Context, u *Session) context.Context {
	newCtx := context.WithValue(ctx, ctxKeyUser, u)
	return newCtx
}

// IsDebugSessionAuth checks whether the context is authenticated by the debug portal.
// The context must be preprocessed by `MWTryAuth`.
func IsDebugSessionAuth(ctx context.Context) bool {
	authBy, err := util.GetContextAuthBy(ctx)
	return err == nil && authBy == authByID
}

// GetDebugSessionUser gets a session user in the context.
// The context must be preprocessed by `MWTryAuth`.
func GetDebugSessionUser(ctx context.Context) (*Session, error) {
	if !IsDebugSessionAuth(ctx) {
		return nil, nil
	}
	i := ctx.Value(ctxKeyUser)
	if i == nil {
		return nil, errcode.NewNotAuthenticatedErr(errors.New("debug session is unauthorized"))
	}
	u, ok := i.(*Session)
	if !ok {
		panic("expect type Session in context")
	}
	return u, nil
}