package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"reflect"
	"time"

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

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

//	{
//	    "code": 206,
//	    "context": ctx,
//	}
type Resumer struct {
	taskDao model.TaskModel
}

func NewResumer(m model.TaskModel) *Resumer {
	return &Resumer{
		taskDao: m,
	}
}

// check event.ResponsePayload["context"] in case dead loop
func (s *Resumer) checkResumeContext(ctx context.Context, event Event) error {
	sched := service.NewLambdaScheduler(s.taskDao)
	inst, err := sched.GetInstance(ctx, event.GetTaskID())
	if err != nil {
		return err
	}
	if reflect.DeepEqual(inst.Data["context"], event.ResponsePayload["context"]) {
		return fmt.Errorf("resume context same with last call: %s", inst.Data["context"])
	}
	return nil
}

func (s *Resumer) ResumeTask(ctx context.Context, event Event) (err error) {
	tasks, err := s.taskDao.GetInstanceList([]string{event.GetTaskID()})
	if err != nil {
		return err
	}
	if len(tasks) == 0 {
		log.Error("task not found", zap.String("id", event.GetTaskID()))
		return fmt.Errorf("no such task id: %s", event.GetTaskID())
	}
	task := &tasks[0]

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

	if !event.NeedResume() {
		runningResult := model.Success
		if event.Failed() {
			runningResult = model.Failed
		}
		log.Info(
			"task finished",
			zap.String("function", event.GetFunctionArn()),
			zap.String("task_id", event.GetTaskID()),
			zap.Any("code", event.ResponsePayload["code"]),
			zap.Any("status_code", event.ResponseContext.StatusCode),
			zap.Int("size", len(event.ResponsePayload)),
		)
		task.Status = runningResult
		task.Data = event.ResponsePayload
		if err := s.taskDao.PutInstance(task); err != nil {
			log.Error("update task status error", zap.Error(err))
			return err
		}

		return s.triggerHook(ctx, event)
	}
	if err := s.checkResumeContext(ctx, event); err != nil {
		return err
	}

	sess := session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))
	svc := lambda.New(sess)

	// Save this call result before next call
	task.Status = model.Running
	task.Data = event.ResponsePayload
	if err := s.taskDao.PutInstance(task); err != nil {
		log.Error("update task status error", zap.Error(err))
		return err
	}
	invokeInput := &lambda.InvokeInput{
		FunctionName:   aws.String(event.GetFunctionArn()),
		InvocationType: aws.String("Event"),
		Payload:        event.GetInvokePayload(),
	}

	result, err := svc.Invoke(invokeInput)
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok {
			log.Error("invoke lambda failed", zap.String("code", aerr.Code()), zap.Error(aerr))
		} else {
			log.Error("invoke lambda failed", zap.Error(aerr))
		}
		return err
	}
	if result.FunctionError != nil {
		log.Error("invoke lambda failed", zap.String("response", result.GoString()))
		return errors.New("invoke lambda failed")
	}
	if result.StatusCode != nil && *result.StatusCode > 204 {
		log.Error("invoke lambda failed", zap.String("response", result.GoString()))
		return errors.New("invoke lambda failed")
	}
	log.Debug("invoke result", zap.String("response", result.GoString()))
	return nil
}

func (s *Resumer) triggerHook(ctx context.Context, event Event) error {
	sched := service.NewLambdaScheduler(s.taskDao)
	inst, err := sched.GetInstance(ctx, event.GetTaskID())
	if err != nil {
		log.Error("get instance error", zap.Error(err))
		return err
	}

	if inst.Hook == "" {
		return nil
	}

	// TODO: clinic will change to whitelist ce-runtime, so the token will
	// be omitted in future versions
	clinicKey := os.Getenv("CLINIC_API_KEY")

	data, err := json.Marshal(inst)
	if err != nil {
		return err
	}
	c := http.Client{
		Timeout: time.Second * 10,
	}
	req, err := http.NewRequest(http.MethodPut, inst.Hook, bytes.NewBuffer(data))
	if err != nil {
		log.Error("build callback request failed", zap.Error(err))
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", clinicKey))
	resp, err := c.Do(req)
	if err != nil {
		log.Error("trigger hook error", zap.Error(err))
		return err
	}
	body, _ := io.ReadAll(resp.Body)
	defer resp.Body.Close()
	log.Info("callback sent", zap.Int("status", resp.StatusCode), zap.ByteString("body", body))
	return nil
}