package main

import (
	"errors"
	"fmt"

	log "github.com/sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws/awserr"
	jwt "github.com/dgrijalva/jwt-go"
	"github.com/lestrrat/go-jwx/jwk"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cognitoidentityprovider"
	"github.com/spf13/viper"
)

// Cognito represents an abstraction of the Amazon Cognito identity provider service.
type Cognito struct {
	cip *cognitoidentityprovider.CognitoIdentityProvider
}

var userPoolID string
var clientID string
var jwksURL string
var keySet *jwk.Set

func init() {

	log.Info("Initializing Cognito")

	log.Info("Loading configuration")
	viper.SetConfigName("config") // config.toml
	viper.AddConfigPath(".")      // use working directory

	if err := viper.ReadInConfig(); err != nil {
		log.Errorf("error reading config file, %v", err)
		return
	}

	userPoolID = viper.GetString("cognito.userPoolID")
	clientID = viper.GetString("cognito.clientID")
	jwksURL = viper.GetString("cognito.jwksURL")

	log.Info("userPoolID: ", userPoolID)
	log.Info("clientID: ", clientID)
	log.Info("jwksURL: ", jwksURL)

	if err := loadKeySet(); err != nil {
		log.Error("Error: ", err)
	}
}

// loadKeySet caches the keyset so we don't have to make a request every time
// we want to verify a JWT
func loadKeySet() error {
	log.Info("Caching keyset")
	var err error
	keySet, err = jwk.FetchHTTP(jwksURL)
	if err != nil {
		return err
	}
	return nil
}

// NewCognito creates a new instance of the Cognito client
func NewCognito() *Cognito {

	c := &Cognito{}

	// Create Session
	sess := session.Must(session.NewSession())
	c.cip = cognitoidentityprovider.New(sess)

	return c
}

// SignUp creates a new Cognito user in the user pool, setting its status
// to CONFIRMED. Returns the authenticated user's JWT access token.
func (c *Cognito) SignUp(username string, password string, email string, fullName string) (string, error) {

	log.Info("AdminCreateUser: ", username)

	// Creates user in FORCE_CHANGE_PASSWORD state
	_, err := c.cip.AdminCreateUser(&cognitoidentityprovider.AdminCreateUserInput{
		Username:          aws.String(username),
		TemporaryPassword: aws.String(password),
		UserPoolId:        aws.String(userPoolID),
		UserAttributes: []*cognitoidentityprovider.AttributeType{
			{
				Name:  aws.String("email_verified"),
				Value: aws.String("true"),
			},
			{
				Name:  aws.String("email"),
				Value: aws.String(email),
			},
			{
				Name:  aws.String("name"),
				Value: aws.String(fullName),
			},
		},
	})

	if err != nil {
		log.Error("Error: ", err.Error())
		return "", err
	}

	// Attempt login to get session value, which is used to confirm the user

	aia := &cognitoidentityprovider.AdminInitiateAuthInput{
		AuthFlow: aws.String("ADMIN_NO_SRP_AUTH"),
		AuthParameters: map[string]*string{
			"USERNAME": aws.String(username),
			"PASSWORD": aws.String(password),
		},
		ClientId:   aws.String(clientID),
		UserPoolId: aws.String(userPoolID),
	}

	log.Info("AdminInitiateAuth: ", username)
	authresp, autherr := c.cip.AdminInitiateAuth(aia)

	log.Info("ChallengeName: ", aws.StringValue(authresp.ChallengeName))

	if autherr != nil {
		log.Warn(autherr.Error())
	}

	// Set user to CONFIRMED

	artaci := &cognitoidentityprovider.AdminRespondToAuthChallengeInput{
		ChallengeName: aws.String("NEW_PASSWORD_REQUIRED"), // Required
		ClientId:      aws.String(clientID),                // Required
		UserPoolId:    aws.String(userPoolID),              // Required
		ChallengeResponses: map[string]*string{
			"USERNAME":     aws.String(username),
			"NEW_PASSWORD": aws.String(password), // Required
		},
		Session: authresp.Session, // session value from AdminInitiateAuth
	}

	log.Info("AdminRespondToAuthChallenge: ", username)
	chalresp, err := c.cip.AdminRespondToAuthChallenge(artaci)

	if err != nil {
		log.Error(err.Error())
		return "", err.(awserr.Error).OrigErr()
	}

	idToken := aws.StringValue(chalresp.AuthenticationResult.IdToken)
	accessToken := aws.StringValue(chalresp.AuthenticationResult.AccessToken)

	log.Debug("ID Token: ", idToken)
	log.Debug("AccessToken: ", accessToken)

	return accessToken, nil
}

// SignIn authenticates a user and returns a JWT token
func (c *Cognito) SignIn(username string, password string) (string, error) {

	aia := &cognitoidentityprovider.AdminInitiateAuthInput{
		AuthFlow: aws.String("ADMIN_NO_SRP_AUTH"),
		AuthParameters: map[string]*string{
			"USERNAME": aws.String(username),
			"PASSWORD": aws.String(password),
		},
		ClientId:   aws.String(clientID),
		UserPoolId: aws.String(userPoolID),
	}

	log.Info("AdminInitiateAuth: ", username)
	authresp, autherr := c.cip.AdminInitiateAuth(aia)

	if autherr != nil {
		log.Error(autherr.Error())
		return "", autherr
	}

	accessToken := aws.StringValue(authresp.AuthenticationResult.AccessToken)

	log.Debug("AccessToken: ", accessToken)

	return accessToken, nil
}

// ValidateToken validates a JWT token and returns the 'sub' claim.
func (c *Cognito) ValidateToken(jwtToken string) (string, error) {

	log.Debug("ValidateToken: ", jwtToken)

	token, err := jwt.Parse(jwtToken, c.getKey)
	if err != nil {
		return "", fmt.Errorf("Could not parse JWT: %v", err)
	}

	log.Debug("JWT signature: ", token.Signature)

	if claims, ok := token.Claims.(jwt.MapClaims); ok && token.Valid {
		if claims["token_use"] != "access" {
			return "", fmt.Errorf("token_use mismatch: %s", claims["token_use"])
		}

		return claims["sub"].(string), nil // Valid token
	}

	return "", nil // Invalid token
}

// getKey returns the key for validating in ValidateToken
func (c *Cognito) getKey(token *jwt.Token) (interface{}, error) {
	keyID, ok := token.Header["kid"].(string)
	if !ok {
		return nil, errors.New("expecting JWT header to have string kid")
	}

	log.Debug("kid: ", keyID)

	if key := keySet.LookupKeyID(keyID); len(key) == 1 {
		return key[0].Materialize()
	}

	return nil, errors.New("unable to find key")
}