package session

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

	"github.com/form3tech-oss/jwt-go"
	"github.com/tidbcloud/debug_portal/auth"
	"github.com/tidbcloud/debug_portal/config"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"

	"github.com/pingcap/errcode"
	"github.com/pingcap/errors"
)

const (
	sessionTimeout                     = time.Hour * 48
	sessionCookieName                  = "debug.session_id"
	authByID                           = "debugSession"
	ctxKeyUser         auth.ContextKey = "DebugSessionToken"
	setHeaderPrefix                    = "x-set-header_"
	queryHeaderPrefix                  = "x-query-header_"
	cookieHeaderPrefix                 = "x-cookie-header_"
	httpMethod                         = "x-http-method"
	httpURI                            = "x-http-uri"
	AuthProcessedByKey auth.ContextKey = "AuthProcessedBy"
)

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

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

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

func NewSessionManager(config *config.SessionConfig) *SessionManager {
	return &SessionManager{sessionSecret: []byte(config.Secret)}
}

// 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(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(), 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) {
	// change sth
	sessionID := FirstHeader(grpcCtx, strings.ToLower(cookieHeaderPrefix+sessionCookieName))
	s, err := sm.getSession(sessionID)
	if err != nil {
		return nil, err
	}

	// CSRF token only need to be checked for non GET requests.
	// change sth
	method := FirstHeader(grpcCtx, httpMethod)
	if method != http.MethodGet {
		csrfToken := 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 wrapContext(ctx context.Context, u *Session) context.Context {
	newCtx := context.WithValue(ctx, ctxKeyUser, u)
	return newCtx
}

// FirstHeader retrives the header content from the context.
// If multiples of the header were given, it gives the first.
func FirstHeader(ctx context.Context, header string) string {
	headers, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		return ""
	}
	headerTexts := headers[strings.ToLower(header)]
	if len(headerTexts) < 1 {
		return ""
	}
	return headerTexts[0]
}

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 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 SetHeader(grpcCtx, "Set-Cookie", cookie.String())
}

func SetHeader(grpcCtx context.Context, name string, value string) error {
	return grpc.SetHeader(grpcCtx, metadata.New(map[string]string{
		setHeaderPrefix + name: value,
	}))
}