// Copyright 2022 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"context"
	"encoding/json"
	"fmt"
	"path/filepath"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/google/uuid"
	"github.com/pingcap/log"
	"go.uber.org/zap"

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

	"github.com/pingcap-inc/compute-engine/pkg/database/model"
)

type Scheduler interface {
	// ListTask returns all lambda functions whose description contains {LambdaDescContains}.
	ListTask(ctx context.Context) ([]string, error)

	// TriggerSync call a lambda function synchronously and waits for result (InvocationType is `RequestResponse`).
	TriggerSync(ctx context.Context, taskName string, tags map[string]interface{}, kwargs map[string]interface{}) (string, map[string]interface{}, error)
	// RetrySync behaves like TriggerSync, but reuse task id
	RetrySync(ctx context.Context, id string) (string, map[string]interface{}, error)

	// TriggerAsync call a lambda function asynchronously (InvocationType is `Event`).
	TriggerAsync(ctx context.Context, taskName string, hook string, tags map[string]interface{}, kwargs map[string]interface{}) (string, error)
	RetryAsync(ctx context.Context, id string) (string, error)

	// A instance is the running status object for a task in the DB.
	// A task could map to multiple lambda function executions, but the instance object is fixed no matter how many times it runs.
	SearchInstance(
		ctx context.Context, taskName string, status model.TaskStatus, begin *time.Time, end *time.Time,
		tags map[string]interface{}, args map[string]interface{}, data map[string]interface{},
	) ([]string, error)
	GetInstance(ctx context.Context, id string) (*model.TaskInstance, error)
	ListInstance(ctx context.Context, ids []string) ([]model.TaskInstance, error)
}

const (
	LambdaTaskRole = "ce-task"
)

// LambdaScheduler implements Scheduler interface
type LambdaScheduler struct {
	dao model.TaskModel
	svc *lambda.Lambda
}

func NewLambdaScheduler(dao model.TaskModel) Scheduler {
	sess := session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))
	svc := lambda.New(sess)
	return &LambdaScheduler{
		dao: dao,
		svc: svc,
	}
}

func (l *LambdaScheduler) checkTask(ctx context.Context, function string) error {
	result, err := l.svc.ListFunctionEventInvokeConfigs(&lambda.ListFunctionEventInvokeConfigsInput{
		FunctionName: aws.String(function),
	})
	if err != nil {
		return err
	}
	for {
		for _, cfg := range result.FunctionEventInvokeConfigs {
			if cfg.DestinationConfig == nil {
				continue
			}
			if cfg.DestinationConfig.OnFailure == nil || cfg.DestinationConfig.OnFailure.Destination == nil {
				continue
			}
			if cfg.DestinationConfig.OnSuccess == nil || cfg.DestinationConfig.OnSuccess.Destination == nil {
				continue
			}
			if !strings.Contains(aws.StringValue(cfg.DestinationConfig.OnSuccess.Destination), "ce-runtime") {
				continue
			}
			if !strings.Contains(aws.StringValue(cfg.DestinationConfig.OnFailure.Destination), "ce-runtime") {
				continue
			}
			return nil
		}
		if result.NextMarker == nil {
			return fmt.Errorf("function %s not set destination to ce-runtime, please check that", function)
		}
		result, err = l.svc.ListFunctionEventInvokeConfigs(&lambda.ListFunctionEventInvokeConfigsInput{
			Marker: result.NextMarker,
		})
		if err != nil {
			return err
		}
	}
}

func (l *LambdaScheduler) mustHasTask(ctx context.Context, function string) error {
	funcs, err := l.ListTask(ctx)
	if err != nil {
		return err
	}
	for _, f := range funcs {
		if f == function {
			return nil
		}
	}
	return fmt.Errorf("function %s not exist", function)
}

func (l *LambdaScheduler) ListTask(ctx context.Context) ([]string, error) {
	var tasks []string

	result, err := l.svc.ListFunctions(nil)
	if err != nil {
		return nil, err
	}
	for {
		for _, f := range result.Functions {
			if filepath.Base(aws.StringValue(f.Role)) != LambdaTaskRole {
				continue
			}
			tasks = append(tasks, aws.StringValue(f.FunctionName))
		}
		if result.NextMarker == nil {
			break
		}
		result, err = l.svc.ListFunctions(&lambda.ListFunctionsInput{
			Marker: result.NextMarker,
		})
		if err != nil {
			return nil, err
		}
	}

	return tasks, nil
}

func (l *LambdaScheduler) triggerSyncWithId(
	ctx context.Context,
	id string,
	taskName string,
	tags map[string]interface{},
	kwargs map[string]interface{},
) (data map[string]interface{}, err error) {
	if taskName != "" {
		if err := l.mustHasTask(ctx, taskName); err != nil {
			return nil, err
		}
	}

	// Reuse the existing task instance.
	task := &model.TaskInstance{ID: id, Task: taskName, Tags: tags, Args: kwargs}
	if tasks, err := l.dao.GetInstanceList([]string{id}); err != nil {
		return nil, err
	} else if len(tasks) > 0 {
		task = &tasks[0]
	}

	// Start the task.
	task.Status = model.Running
	if err := l.dao.PutInstance(task); err != nil {
		return nil, err
	}
	// Task status update hook.
	defer func() {
		if err := l.dao.PutInstance(task); err != nil {
			log.Error("update task status error", zap.Error(err))
		}
	}()

	// Build lambda input.
	payload, err := json.Marshal(job.InputPayload{
		TaskID: task.ID,
		Tags:   task.Tags,
		Kwargs: task.Args,
	})
	if err != nil {
		task.Status = model.Failed
		return nil, err
	}

	// Call lambda function.
	input := &lambda.InvokeInput{
		FunctionName: aws.String(task.Task),
		Payload:      payload,
	}

	// Handle lambda service or sdk error.
	output, err := l.svc.Invoke(input)
	if err != nil {
		task.Status = model.Failed
		return nil, err
	}

	task.Data = map[string]any{}
	if err := json.Unmarshal(output.Payload, &task.Data); err != nil {
		task.Status = model.Failed
		return nil, err
	}

	// Handle lambda function error.
	task.Status = model.Success
	if aws.Int64Value(output.StatusCode) != 200 || output.FunctionError != nil || task.Data["errorType"] != nil || task.Data["errorMessage"] != nil {
		task.Status = model.Failed
		return nil, fmt.Errorf("invoking lambda function %s returns status code %d, payload: %s, error: %s",
			task.Task, aws.Int64Value(output.StatusCode), string(output.Payload), aws.StringValue(output.FunctionError))
	}

	return task.Data, nil
}

func (l *LambdaScheduler) TriggerSync(
	ctx context.Context,
	taskName string,
	tags map[string]interface{},
	kwargs map[string]interface{},
) (string, map[string]interface{}, error) {
	id := uuid.New().String()
	data, err := l.triggerSyncWithId(ctx, id, taskName, tags, kwargs)
	return id, data, err
}

func (l *LambdaScheduler) RetrySync(ctx context.Context, id string) (string, map[string]interface{}, error) {
	data, err := l.triggerSyncWithId(ctx, id, "", nil, nil)
	return id, data, err
}

func (l *LambdaScheduler) triggerAsyncWithId(
	ctx context.Context,
	id string,
	taskName string,
	hook string,
	tags map[string]interface{},
	kwargs map[string]interface{},
) (err error) {
	if taskName != "" {
		if err := l.mustHasTask(ctx, taskName); err != nil {
			return err
		}
		if err := l.checkTask(ctx, taskName); err != nil {
			return err
		}
	}

	task := &model.TaskInstance{ID: id, Task: taskName, Hook: hook, Tags: tags, Args: kwargs}
	if tasks, err := l.dao.GetInstanceList([]string{id}); err != nil {
		return err
	} else if len(tasks) > 0 {
		task = &tasks[0]
	}

	task.Status = model.Running
	if err := l.dao.PutInstance(task); err != nil {
		return err
	}
	defer func() {
		if err != nil {
			task.Status = model.Failed
			if err := l.dao.PutInstance(task); err != nil {
				log.Error("update task status error", zap.Error(err))
			}
		}
	}()

	payload, err := json.Marshal(job.InputPayload{
		TaskID: task.ID,
		Tags:   task.Tags,
		Kwargs: task.Args,
	})
	if err != nil {
		return err
	}
	input := &lambda.InvokeInput{
		FunctionName:   aws.String(task.Task),
		InvocationType: aws.String("Event"),
		Payload:        payload,
	}
	output, err := l.svc.Invoke(input)
	if err != nil {
		return err
	}
	if aws.Int64Value(output.StatusCode) != 202 || output.FunctionError != nil {
		return fmt.Errorf("invoking lambda function %s returns status code %d, payload: %s, error: %s",
			task.Task, aws.Int64Value(output.StatusCode), string(output.Payload), aws.StringValue(output.FunctionError))
	}

	return nil
}

func (l *LambdaScheduler) TriggerAsync(
	ctx context.Context,
	taskName string,
	hook string,
	tags map[string]interface{},
	kwargs map[string]interface{},
) (string, error) {
	id := uuid.New().String()
	return id, l.triggerAsyncWithId(ctx, id, taskName, hook, tags, kwargs)
}

func (l *LambdaScheduler) RetryAsync(ctx context.Context, id string) (string, error) {
	return id, l.triggerAsyncWithId(ctx, id, "", "", nil, nil)
}

func (l *LambdaScheduler) SearchInstance(
	ctx context.Context,
	taskName string,
	status model.TaskStatus,
	begin *time.Time,
	end *time.Time,
	tags map[string]interface{},
	args map[string]interface{},
	data map[string]interface{},
) ([]string, error) {
	ids := []string{}
	for _, s := range []model.TaskStatus{model.Pending, model.Running, model.Success, model.Failed} {
		if status&s == 0 {
			continue
		}
		xs, err := l.dao.SearchInstance(taskName, int(s), begin, end, tags, args, data)
		if err != nil {
			return nil, err
		}
		ids = append(ids, xs...)
	}
	return ids, nil
}

func (l *LambdaScheduler) GetInstance(ctx context.Context, id string) (*model.TaskInstance, error) {
	insts, err := l.dao.GetInstanceList([]string{id})
	if err != nil {
		return nil, err
	}
	if len(insts) == 0 {
		return nil, fmt.Errorf("no instance with id %s found", id)
	}
	return &insts[0], nil
}

func (l *LambdaScheduler) ListInstance(ctx context.Context, ids []string) ([]model.TaskInstance, error) {
	return l.dao.GetInstanceList(ids)
}