package main

import (
	"context"
	"encoding/json"
	"log"
	"os"
	"strconv"
	"strings"

	runtime "github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-lambda-go/lambdacontext"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/secretsmanager"
	"github.com/samber/lo"

	"github.com/pingcap-inc/compute-engine/pkg/job"
	"github.com/pingcap-inc/compute-engine/pkg/unstable_plan"
)

const (
	defaultOffsetHours = 24
	offsetHoursEnvKey  = "OFFSET_HOURS"

	defaultSecretRegion = "us-west-2"
	secretRegionEnvKey  = "SECRET_REGION"
)

func handleRequest(ctx context.Context, payload job.InputPayload) (RespWrapper, error) {
	// Log event.
	eventJson, _ := json.MarshalIndent(payload, "", "  ")
	log.Printf("EVENT: %s", eventJson)
	// Log environment variables.
	log.Printf("REGION: %s", os.Getenv("AWS_REGION"))
	log.Println("ALL ENV VARS:")
	for _, element := range os.Environ() {
		log.Println(element)
	}
	// Log request context.
	lc, _ := lambdacontext.FromContext(ctx)
	log.Printf("REQUEST ID: %s", lc.AwsRequestID)
	log.Printf("FUNCTION NAME: %s", lambdacontext.FunctionName)
	deadline, _ := ctx.Deadline()
	log.Printf("DEADLINE: %s", deadline)

	// Get offset hours from environment variable.
	offsetHours, err := strconv.ParseUint(os.Getenv(offsetHoursEnvKey), 10, 64)
	if err != nil {
		log.Printf("failed to get offset hours from environment variable %s, use default value %d instead", offsetHoursEnvKey, defaultOffsetHours)
		offsetHours = defaultOffsetHours
	}

	// Get secret region from environment variable.
	secretRegion := os.Getenv(secretRegionEnvKey)
	if secretRegion == "" {
		log.Printf("failed to get secret region from environment variable %s, use default value %s instead", secretRegionEnvKey, defaultSecretRegion)
		secretRegion = defaultSecretRegion
	}

	// Extract input parameters.
	dbAddr := payload.Kwargs["addr"].(string)
	dbPort := int(payload.Kwargs["port"].(float64))
	dbUser := payload.Kwargs["user"].(string)
	dbPasswdSecret := payload.Kwargs["passwdSecret"].(string)
	taskCreateTimeEpoch := int64(payload.Kwargs["taskCreateTimeEpoch"].(float64))
	onlyValidateAccount := payload.Kwargs["onlyValidateAccount"].(bool)
	if offsetHoursFromRequest, ok := payload.Kwargs["offsetHours"]; ok {
		offsetHours = uint64(offsetHoursFromRequest.(float64))
	}

	// Get password from secret manager.
	dbPasswd, err := getPasswdFromSecretManager(dbPasswdSecret, secretRegion)
	if err != nil {
		log.Printf("failed to get password from secret manager: %v", err)
		return RespWrapper{}, err
	}

	// Init customer TiDB connection.
	db, err := unstable_plan.NewGormDB(dbUser, dbPasswd, dbAddr, dbPort)
	if err != nil {
		log.Printf("failed to init customer TiDB connection: %v", err)
		return RespWrapper{}, err
	}

	// Create task execution context service.
	svc := unstable_plan.NewService(db, taskCreateTimeEpoch, int(offsetHours))

	if onlyValidateAccount {
		m, err := svc.CheckPrivilege()
		if err != nil {
			log.Printf("failed to check privilege: %v", err)
			return RespWrapper{}, err
		}
		return RespWrapper{Data: []unstable_plan.Model{*m}}, nil
	}

	res, err := svc.DoComputation()
	if err != nil {
		log.Printf("failed to execute task: %v", err)
		return RespWrapper{}, err
	}

	return RespWrapper{Data: res}, nil
}

func getPasswdFromSecretManager(secretName, region string) (string, error) {
	awsConfig := aws.NewConfig().WithMaxRetries(3).WithRegion(region)
	opts := session.Options{Config: *awsConfig}
	sess, err := session.NewSessionWithOptions(opts)
	if err != nil {
		return "", err
	}
	svc := secretsmanager.New(sess)
	input := &secretsmanager.GetSecretValueInput{
		SecretId:     aws.String(secretName),
		VersionStage: aws.String("AWSCURRENT"),
	}
	result, err := svc.GetSecretValue(input)
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(lo.FromPtr(result.SecretString)), nil
}

type RespWrapper struct {
	Data []unstable_plan.Model `json:"data"`
}

func main() {
	runtime.Start(handleRequest)
}