package middleware

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

	"code.crute.us/mcrute/cloud-identity-broker/app/models"
	"code.crute.us/mcrute/cloud-identity-broker/auth"
	"code.crute.us/mcrute/cloud-identity-broker/auth/github"

	glecho "code.crute.us/mcrute/golib/echo"
	"github.com/labstack/echo/v4"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"golang.org/x/oauth2"
)

// apiKeyRequests tracks the number of requests made with the legacy X-API-Key
// header instead of the Authorization header so that we can track when it's
// safe to remove that edge-case.
var apiKeyRequests = promauto.NewCounterVec(prometheus.CounterOpts{
	Namespace: "aws_access", // Legacy Namespace
	Name:      "api_key_request_total",
	Help:      "Total number of requests using the X-API-Key header",
}, nil)

const (
	authPrincipalContextKey = "broker.AuthorizedPrincipal"
	gitHubTokenCookie       = "github-token"
	gitHubStateCookie       = "github-state"
	oauthReturnUrl          = "/github-auth"
)

// GetAuthorizedPrincipal returns the user principal object from the request
// context and casts it correctly. Will return error if there is no principal
// or if the principal is of the incorrect type.
//
// Note that use of this function implies that AuthenticationMiddleware is used
// somewhere in the stack before the handler calling this function is
// dispatched.
func GetAuthorizedPrincipal(c echo.Context) (*models.User, error) {
	rp := c.Get(authPrincipalContextKey)
	if rp == nil {
		return nil, fmt.Errorf("No principal set in request")
	}
	principal, ok := rp.(*models.User)
	if !ok {
		return nil, fmt.Errorf("Principal in request is not of User type")
	}
	return principal, nil
}

type AuthenticationMiddleware struct {
	Store          models.UserStore
	JWTManager     *auth.JWTManager
	GitHub         *github.GitHubAuthenticator
	CookieDuration time.Duration
}

func (m *AuthenticationMiddleware) redirectToGitHubAuth(c echo.Context) error {
	redir, state := m.GitHub.GetAuthRedirect()

	c.SetCookie(&http.Cookie{
		Name:     gitHubStateCookie,
		Value:    state,
		Path:     "/",
		Secure:   true,
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
	})

	return c.Redirect(http.StatusFound, redir)
}

// RegisterUrls registers the URLs required by this middleware and handler with an echo instance.
//
// This is here instead of in the web main because these paths are encoded in
// the configuration for the GitHub application so changing them requires
// addition changes to that configuration.
func (m *AuthenticationMiddleware) RegisterUrls(e glecho.URLRouter) {
	e.GET(oauthReturnUrl, m.HandleCompleteLogin)
}

// Middleware does user authentication based on either an X-API-Key header,
// Authorization header, or GitHub cookie depending on how the request is
// phrased.
//
// If the request has either an X-API-Key or an Authorization Bearer header
// then that must pass validation with the downstream validation logic.
// Failures through this path are hard failures and the only way to re-try them
// is to authenticate with a new token. The underlying assumption is that only
// programmatic access goes through this path so redirecting to interactive
// authentication is pointless.
//
// In the absence of those headers it's assumed that the user is interactive
// and their auth cookie will be read and validated (by the exact same logic
// that an API key is validated, they're the same format) but the failure case
// here will redirect the user to GitHub for interactive auth.
//
// X-API-Key should be considered deprecated and the Authorization header with
// a type of Bearer should be used instead. This is more in-line with Oauth 2
// style authentication. However, for now this middleware continues to support
// X-API-Key for to not break legacy API clients.
func (m *AuthenticationMiddleware) Middleware(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		token := c.Request().Header.Get("X-API-Key")
		if token == "" {
			tp := strings.Split(c.Request().Header.Get(echo.HeaderAuthorization), " ")
			if len(tp) == 2 && tp[0] == "Bearer" {
				token = tp[1]
			}
		} else {
			apiKeyRequests.With(nil).Inc()
		}

		// If an API key is specified this is a success or failure path. There
		// is no option to authenticate to GitHub as would an interactive user.
		if token != "" {
			u, err := m.JWTManager.Validate(token)
			if err != nil {
				c.Logger().Debugf("Error validating JWT: %w", err)
				return echo.ErrUnauthorized
			}
			c.Set(authPrincipalContextKey, u)
			return next(c)
		}

		// Kick them to GitHub auth if they have no auth cookie
		authCookie, err := c.Cookie(gitHubTokenCookie)
		if err != nil {
			return m.redirectToGitHubAuth(c)
		}

		// If they fail the check them bounce them through logout to remove
		// their existing cookies which should then bounce them back through
		// GitHub auth, which will eventually land them back here.
		u, err := m.JWTManager.Validate(authCookie.Value)
		if err != nil {
			c.Logger().Debugf("Error validating JWT: %w", err)
			return c.Redirect(http.StatusFound, "/logout")
		}

		c.Set(authPrincipalContextKey, u)
		return next(c)
	}
}

// HandleCompleteLogin handles the Oauth 2 code flow. It receives the auth code
// and uses that to retrieve the auth token. This sets the user's auth cookie
// to a authenticated JWT.
//
// This is redirected-to by the Oauth authorization server and should never be
// hit directly by a user or script.
func (m *AuthenticationMiddleware) HandleCompleteLogin(c echo.Context) error {
	ctx := context.Background()

	code, state := c.QueryParam("code"), c.QueryParam("state")
	if code == "" || state == "" {
		return echo.ErrBadRequest
	}

	ghState, err := c.Cookie(gitHubStateCookie)
	if err != nil || ghState.Value == "" {
		return echo.ErrBadRequest
	}

	if ghState.Value != state {
		return echo.ErrBadRequest
	}

	token, err := m.GitHub.GetTokens(code)
	if err != nil {
		return echo.ErrBadRequest
	}

	user, err := m.GitHub.GetUsernameWithToken(token.AccessToken)
	if err != nil {
		c.Logger().Debugf("Error getting GitHub username with token: %w", err)
		return echo.ErrUnauthorized
	}

	dbUser, err := m.Store.Get(ctx, user)
	if err != nil {
		c.Logger().Errorf("GitHub user %s does not have access to app", user)
		return echo.ErrUnauthorized
	}

	// Service users should only be allowed to submit self-signed JWTs. A
	// service user should never be able to use GitHub auth.
	if dbUser.IsService {
		c.Logger().Errorf("Service user %s attempted to use GitHub auth", user)
		return echo.ErrUnauthorized
	}

	jwt, sk, err := m.JWTManager.CreateForUser(dbUser)
	if err != nil {
		return echo.ErrInternalServerError
	}

	dbUser.AddKey(sk)
	dbUser.GCKeys() // This is a convenient place to do it

	dbUser.AddToken("github", &oauth2.Token{
		AccessToken:  token.AccessToken,
		RefreshToken: token.RefreshToken,
	})

	if err := m.Store.Put(ctx, dbUser); err != nil {
		return echo.ErrInternalServerError
	}

	c.SetCookie(&http.Cookie{
		Name:     gitHubTokenCookie,
		Value:    jwt,
		Path:     "/",
		MaxAge:   int(m.CookieDuration.Seconds()),
		Secure:   true,
		SameSite: http.SameSiteLaxMode,
	})

	return c.Redirect(http.StatusFound, "/")
}