// 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 dynamo

import (
	"encoding/json"
	"reflect"
	"sort"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbattribute"
	"github.com/aws/aws-sdk-go/service/dynamodb/expression"
	"github.com/pingcap/log"
	"go.uber.org/zap"

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

const (
	TaskTableName = "compute-engine-task-instances"
	TagsTableName = "compute-engine-task-tags"
	ArgsTableName = "compute-engine-task-args"
	DataTableName = "compute-engine-task-data"

	TaskTableSearchIndex = "task_name-created_at-index"
)

type DynamoDB interface {
	TransactWriteItems(input *dynamodb.TransactWriteItemsInput) (*dynamodb.TransactWriteItemsOutput, error)
	BatchGetItem(input *dynamodb.BatchGetItemInput) (*dynamodb.BatchGetItemOutput, error)
	Query(input *dynamodb.QueryInput) (*dynamodb.QueryOutput, error)
}

type TaskDao struct {
	db DynamoDB
}

func NewModel() (model.TaskModel, error) {
	sess := session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))
	svc := dynamodb.New(sess)
	return &TaskDao{svc}, nil
}

// PutInstance create or update instance via primary key: ID
func (dao *TaskDao) PutInstance(instance *model.TaskInstance) error {
	var prev *Task
	if insts, err := dao.getInstanceList([]string{instance.ID}, true); err != nil {
		return err
	} else if len(insts) > 0 {
		prev = &insts[0]
	}

	var taglist []KeyValuePair
	for k, v := range instance.Tags {
		bytes, err := json.Marshal(v)
		if err != nil {
			return err
		}
		tag := KeyValuePair{
			TaskID: instance.ID,
			Key:    k,
			Value:  string(bytes),
		}
		taglist = append(taglist, tag)
	}
	var arglist []KeyValuePair
	for k, v := range instance.Args {
		bytes, err := json.Marshal(v)
		if err != nil {
			return err
		}
		arg := KeyValuePair{
			TaskID: instance.ID,
			Key:    k,
			Value:  string(bytes),
		}
		arglist = append(arglist, arg)
	}
	var datalist []KeyValuePair
	for k, v := range instance.Data {
		bytes, err := json.Marshal(v)
		if err != nil {
			return err
		}
		data := KeyValuePair{
			TaskID: instance.ID,
			Key:    k,
			Value:  string(bytes),
		}
		datalist = append(datalist, data)
	}

	update_ts := time.Now().Unix()
	create_ts := update_ts
	if prev != nil {
		create_ts = prev.CreatedAt
	}
	task := &Task{
		TaskID:    instance.ID,
		TaskName:  instance.Task,
		Status:    int(instance.Status),
		Hook:      instance.Hook,
		CreatedAt: create_ts,
		UpdatedAt: update_ts,
		Tags:      taglist,
		Args:      arglist,
		Data:      datalist,
	}

	input, err := dao.dynamoInput(task, prev)
	if err != nil {
		return err
	}
	if _, err := dao.db.TransactWriteItems(input); err != nil {
		return err
	}
	return nil
}

// GetInstanceList query instance via ids
func (dao *TaskDao) GetInstanceList(ids []string) ([]model.TaskInstance, error) {
	insts, err := dao.getInstanceList(ids, true)
	if err != nil {
		log.Error("get instance list error", zap.Error(err))
		return nil, err
	}
	sort.Slice(insts, func(i, j int) bool {
		return insts[i].CreatedAt < insts[j].CreatedAt
	})
	var instances []model.TaskInstance
	for _, inst := range insts {
		i, err := inst.Instance()
		if err != nil {
			return nil, err
		}
		instances = append(instances, *i)
	}
	return instances, nil
}

// SearchInstance query instance in task partition, filter by status, [begin, end)
// The tags, args and data only be filtered by the top level key-values
func (dao *TaskDao) SearchInstance(
	taskName string,
	status int,
	begin *time.Time,
	end *time.Time,
	tags map[string]interface{},
	args map[string]interface{},
	data map[string]interface{},
) ([]string, error) {
	keyCondition := expression.Key("task_name").Equal(expression.Value(taskName))
	filter := expression.Name("status").Equal(expression.Value(status))
	if begin != nil {
		if end != nil {
			keyCondition = keyCondition.And(expression.Key("created_at").Between(expression.Value(begin.Unix()), expression.Value(end.Unix()-1)))
		} else {
			keyCondition = keyCondition.And(expression.Key("created_at").GreaterThanEqual(expression.Value(begin.Unix())))
		}
	} else if end != nil {
		keyCondition = keyCondition.And(expression.Key("created_at").LessThan(expression.Value(end.Unix())))
	}

	expr, err := expression.NewBuilder().WithKeyCondition(keyCondition).WithFilter(filter).Build()
	if err != nil {
		log.Error("build expression error", zap.Error(err))
		return nil, err
	}
	params := &dynamodb.QueryInput{
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		KeyConditionExpression:    expr.KeyCondition(),
		FilterExpression:          expr.Filter(),
		ProjectionExpression:      expr.Projection(),
		TableName:                 aws.String(TaskTableName),
		IndexName:                 aws.String(TaskTableSearchIndex),
	}

	ids, err := dao.filterInstance(params)
	if err != nil {
		return nil, err
	}
	ids, err = dao.filterAttribute(TagsTableName, ids, tags)
	if err != nil {
		return nil, err
	}
	ids, err = dao.filterAttribute(ArgsTableName, ids, args)
	if err != nil {
		return nil, err
	}
	ids, err = dao.filterAttribute(DataTableName, ids, data)
	if err != nil {
		return nil, err
	}
	return ids, nil
}

// No need to close dynamodb connections because there is not one
func (dao *TaskDao) Close() error {
	return nil
}

func (dao *TaskDao) query(input *dynamodb.QueryInput) ([]map[string]*dynamodb.AttributeValue, error) {
	var items []map[string]*dynamodb.AttributeValue
	result, err := dao.db.Query(input)
	if err != nil {
		log.Error("query table error", zap.String("table", aws.StringValue(input.TableName)), zap.Error(err))
		return nil, err
	}
	items = append(items, result.Items...)

	for result.LastEvaluatedKey != nil {
		input = input.SetExclusiveStartKey(result.LastEvaluatedKey)
		result, err = dao.db.Query(input)
		if err != nil {
			log.Error("query table error", zap.String("table", aws.StringValue(input.TableName)), zap.Error(err))
			return nil, err
		}
		items = append(items, result.Items...)
	}
	return items, nil
}

func (dao *TaskDao) get(input *dynamodb.BatchGetItemInput) (map[string][]map[string]*dynamodb.AttributeValue, error) {
	items := map[string][]map[string]*dynamodb.AttributeValue{}

	result, err := dao.db.BatchGetItem(input)
	if err != nil {
		log.Error("batch get error", zap.Error(err))
		return nil, err
	}
	for table, attrs := range result.Responses {
		if items[table] == nil {
			items[table] = []map[string]*dynamodb.AttributeValue{}
		}
		items[table] = append(items[table], attrs...)
	}

	if len(result.UnprocessedKeys) > 0 {
		input.SetRequestItems(result.UnprocessedKeys)
		result, err = dao.db.BatchGetItem(input)
		if err != nil {
			log.Error("batch get error", zap.Error(err))
			return nil, err
		}
		for table, attrs := range result.Responses {
			if items[table] == nil {
				items[table] = []map[string]*dynamodb.AttributeValue{}
			}
			items[table] = append(items[table], attrs...)
		}
	}

	return items, nil
}

func (dao *TaskDao) getAttribute(tbname string, ids, keys []string) (map[string][]KeyValuePair, error) {
	atrributeMap := make(map[string][]KeyValuePair)

	var ks []map[string]*dynamodb.AttributeValue
	for _, id := range ids {
		for _, key := range keys {
			ks = append(ks, map[string]*dynamodb.AttributeValue{
				"task_id": {
					S: aws.String(id),
				},
				"key": {
					S: aws.String(key),
				},
			})
		}
	}

	for _, ks := range utils.Chunk[map[string]*dynamodb.AttributeValue](ks, 100) {
		input := &dynamodb.BatchGetItemInput{
			RequestItems: map[string]*dynamodb.KeysAndAttributes{
				tbname: {
					Keys: ks,
				},
			},
		}
		m, err := dao.get(input)
		if err != nil {
			return nil, err
		}

		for _, item := range m[tbname] {
			data := KeyValuePair{}
			if err := dynamodbattribute.UnmarshalMap(item, &data); err != nil {
				log.Error("unmarshal dynamodb result error", zap.Error(err))
				return nil, err
			}
			if atrributeMap[data.TaskID] == nil {
				atrributeMap[data.TaskID] = make([]KeyValuePair, 0)
			}
			atrributeMap[data.TaskID] = append(atrributeMap[data.TaskID], data)
		}
	}

	return atrributeMap, nil
}

func (dao *TaskDao) loadAttribute(tbname string, ids, keys []string) (map[string][]KeyValuePair, error) {
	atrributeMap := make(map[string][]KeyValuePair)
	if len(ids) == 0 {
		return atrributeMap, nil
	}

	if len(keys) > 0 {
		return dao.getAttribute(tbname, ids, keys)
	}

	for _, id := range ids {
		keyCondition := expression.Key("task_id").Equal(expression.Value(id))
		expr, err := expression.NewBuilder().WithKeyCondition(keyCondition).Build()
		if err != nil {
			log.Error("build expression error", zap.Error(err))
			return nil, err
		}

		params := &dynamodb.QueryInput{
			ExpressionAttributeNames:  expr.Names(),
			ExpressionAttributeValues: expr.Values(),
			KeyConditionExpression:    expr.KeyCondition(),
			FilterExpression:          expr.Filter(),
			ProjectionExpression:      expr.Projection(),
			TableName:                 aws.String(tbname),
		}

		items, err := dao.query(params)
		if err != nil {
			return nil, err
		}
		for _, item := range items {
			data := KeyValuePair{}
			if err := dynamodbattribute.UnmarshalMap(item, &data); err != nil {
				log.Error("unmarshal dynamodb result error", zap.Error(err))
				return nil, err
			}
			if atrributeMap[data.TaskID] == nil {
				atrributeMap[data.TaskID] = make([]KeyValuePair, 0)
			}
			atrributeMap[data.TaskID] = append(atrributeMap[data.TaskID], data)
		}
	}

	return atrributeMap, nil
}

func (dao *TaskDao) getInstanceList(ids []string, preload bool) ([]Task, error) {
	if len(ids) == 0 {
		return nil, nil
	}
	var instances []Task

	for _, ids := range utils.Chunk[string](ids, 100) {
		input := &dynamodb.BatchGetItemInput{
			RequestItems: map[string]*dynamodb.KeysAndAttributes{
				TaskTableName: {
					Keys: []map[string]*dynamodb.AttributeValue{},
				},
			},
		}
		for _, id := range ids {
			input.RequestItems[TaskTableName].Keys = append(
				input.RequestItems[TaskTableName].Keys,
				map[string]*dynamodb.AttributeValue{
					"id": {
						S: aws.String(id),
					},
				})
		}
		m, err := dao.get(input)
		if err != nil {
			log.Error("get instance list error", zap.Error(err))
		}
		for _, item := range m[TaskTableName] {
			instance := Task{}
			if err := dynamodbattribute.UnmarshalMap(item, &instance); err != nil {
				log.Error("unmarshal dynamodb result error", zap.Error(err))
				return nil, err
			}
			instances = append(instances, instance)
		}
	}

	if !preload {
		return instances, nil
	}

	tags, err := dao.loadAttribute(TagsTableName, ids, nil)
	if err != nil {
		log.Error("load tags error", zap.Error(err))
		return nil, err
	}
	args, err := dao.loadAttribute(ArgsTableName, ids, nil)
	if err != nil {
		log.Error("load args error", zap.Error(err))
		return nil, err
	}
	data, err := dao.loadAttribute(DataTableName, ids, nil)
	if err != nil {
		log.Error("load data error", zap.Error(err))
		return nil, err
	}

	for idx := range instances {
		id := instances[idx].TaskID
		instances[idx].Tags = tags[id]
		instances[idx].Args = args[id]
		instances[idx].Data = data[id]
	}
	return instances, nil
}

func (dao *TaskDao) diffAttribute(tb string, current []KeyValuePair, prev []KeyValuePair) ([]*dynamodb.TransactWriteItem, error) {
	var transactions []*dynamodb.TransactWriteItem

	cmap := map[string]string{}
	pmap := map[string]string{}

	for _, attr := range current {
		cmap[attr.TaskID+attr.Key] = attr.Value
	}
	for _, attr := range prev {
		pmap[attr.TaskID+attr.Key] = attr.Value
	}

	for _, attr := range current {
		pv, ok := pmap[attr.TaskID+attr.Key]
		if !ok || pv != attr.Value {
			// should add or modify
			xav, err := dynamodbattribute.MarshalMap(attr)
			if err != nil {
				log.Error("marshal error", zap.Error(err))
				return nil, err
			}
			transactions = append(transactions, &dynamodb.TransactWriteItem{
				Put: &dynamodb.Put{
					TableName: aws.String(tb),
					Item:      xav,
				},
			})
		}
	}
	for _, attr := range prev {
		_, ok := cmap[attr.TaskID+attr.Key]
		if !ok {
			// should delete
			xav, err := dynamodbattribute.MarshalMap(attr)
			if err != nil {
				log.Error("marshal error", zap.Error(err))
				return nil, err
			}
			transactions = append(transactions, &dynamodb.TransactWriteItem{
				Delete: &dynamodb.Delete{
					TableName: aws.String(tb),
					Key: map[string]*dynamodb.AttributeValue{
						"task_id": xav["task_id"],
						"key":     xav["key"],
					},
				},
			})
		}
	}

	return transactions, nil
}

func (dao *TaskDao) dynamoInput(current *Task, prev *Task) (*dynamodb.TransactWriteItemsInput, error) {
	var transactions []*dynamodb.TransactWriteItem

	task, err := dynamodbattribute.MarshalMap(current)
	if err != nil {
		log.Error("marshal task error", zap.Error(err))
		return nil, err
	}
	transactions = append(transactions, &dynamodb.TransactWriteItem{
		Put: &dynamodb.Put{
			TableName: aws.String(TaskTableName),
			Item:      task,
		},
	})

	if prev == nil {
		// A pure new task
		for tb, xs := range map[string][]KeyValuePair{
			TagsTableName: current.Tags,
			ArgsTableName: current.Args,
			DataTableName: current.Data,
		} {
			for _, x := range xs {
				xav, err := dynamodbattribute.MarshalMap(x)
				if err != nil {
					log.Error("marshal error", zap.Error(err))
					return nil, err
				}
				transactions = append(transactions, &dynamodb.TransactWriteItem{
					Put: &dynamodb.Put{
						TableName: aws.String(tb),
						Item:      xav,
					},
				})
			}
		}
	} else {
		// should do update
		txns, err := dao.diffAttribute(TagsTableName, current.Tags, prev.Tags)
		if err != nil {
			log.Error("diff tags error", zap.Error(err))
			return nil, err
		}
		transactions = append(transactions, txns...)

		txns, err = dao.diffAttribute(ArgsTableName, current.Args, prev.Args)
		if err != nil {
			log.Error("diff args error", zap.Error(err))
			return nil, err
		}
		transactions = append(transactions, txns...)

		txns, err = dao.diffAttribute(DataTableName, current.Data, prev.Data)
		if err != nil {
			log.Error("diff args error", zap.Error(err))
			return nil, err
		}
		transactions = append(transactions, txns...)
	}

	return &dynamodb.TransactWriteItemsInput{
		TransactItems: transactions,
	}, nil
}

func (dao *TaskDao) filterInstance(params *dynamodb.QueryInput) ([]string, error) {
	items, err := dao.query(params)
	if err != nil {
		return nil, err
	}
	var ids []string
	for _, item := range items {
		instance := Task{}
		if err := dynamodbattribute.UnmarshalMap(item, &instance); err != nil {
			log.Error("unmarshal dynamodb result error", zap.Error(err))
			return nil, err
		}
		ids = append(ids, instance.TaskID)
	}
	return ids, nil
}

func (dao *TaskDao) filterAttribute(tb string, ids []string, attrs map[string]interface{}) ([]string, error) {
	if len(ids) == 0 || len(attrs) == 0 {
		return ids, nil
	}

	var keys []string
	kv := map[string]string{}
	for k, v := range attrs {
		vv, err := json.Marshal(v)
		if err != nil {
			return nil, err
		}
		kv[k] = string(vv)
		keys = append(keys, k)
	}

	attrList, err := dao.loadAttribute(tb, ids, keys)
	if err != nil {
		return nil, err
	}

	ids = []string{}
	for id, xs := range attrList {
		akv := map[string]string{}
		for _, attr := range xs {
			akv[attr.Key] = attr.Value
		}

		if reflect.DeepEqual(kv, akv) {
			ids = append(ids, id)
		}
	}

	return ids, nil
}