// 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 (
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbattribute"
	"github.com/stretchr/testify/assert"

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

type MockDynamoDB struct {
	db map[string]map[string]map[string]*dynamodb.AttributeValue
}

func NewMockDynamoDB() *MockDynamoDB {
	return &MockDynamoDB{
		db: make(map[string]map[string]map[string]*dynamodb.AttributeValue),
	}
}

func (m *MockDynamoDB) key(tbname string, item map[string]*dynamodb.AttributeValue) (string, error) {
	if tbname == TaskTableName {
		task := Task{}
		if err := dynamodbattribute.UnmarshalMap(item, &task); err != nil {
			return "", err
		}
		return task.TaskID, nil
	} else {
		pair := KeyValuePair{}
		if err := dynamodbattribute.UnmarshalMap(item, &pair); err != nil {
			return "", err
		}
		return pair.TaskID + ":" + pair.Key, nil
	}
}

func (m *MockDynamoDB) TransactWriteItems(input *dynamodb.TransactWriteItemsInput) (*dynamodb.TransactWriteItemsOutput, error) {
	for _, txn := range input.TransactItems {
		if txn.Put != nil {
			key, err := m.key(aws.StringValue(txn.Put.TableName), txn.Put.Item)
			if err != nil {
				return nil, err
			}
			if m.db[aws.StringValue(txn.Put.TableName)] == nil {
				m.db[aws.StringValue(txn.Put.TableName)] = make(map[string]map[string]*dynamodb.AttributeValue)
			}
			m.db[aws.StringValue(txn.Put.TableName)][key] = txn.Put.Item
		} else if txn.Delete != nil {
			key, err := m.key(aws.StringValue(txn.Delete.TableName), txn.Delete.Key)
			if err != nil {
				return nil, err
			}
			if m.db[aws.StringValue(txn.Delete.TableName)] == nil {
				m.db[aws.StringValue(txn.Delete.TableName)] = make(map[string]map[string]*dynamodb.AttributeValue)
			}
			delete(m.db[aws.StringValue(txn.Delete.TableName)], key)
		}
	}
	return nil, nil
}

func (m *MockDynamoDB) BatchGetItem(input *dynamodb.BatchGetItemInput) (*dynamodb.BatchGetItemOutput, error) {
	responses := make(map[string][]map[string]*dynamodb.AttributeValue)
	for table, item := range input.RequestItems {
		for _, k := range item.Keys {
			key, err := m.key(table, k)
			if err != nil {
				return nil, err
			}
			if m.db[table][key] != nil {
				responses[table] = append(responses[table], m.db[table][key])
			}
		}
	}

	return &dynamodb.BatchGetItemOutput{
		Responses: responses,
	}, nil
}

// Query for mock dynamodb returns all items
func (m *MockDynamoDB) Query(input *dynamodb.QueryInput) (*dynamodb.QueryOutput, error) {
	items := []map[string]*dynamodb.AttributeValue{}
	for _, item := range m.db[aws.StringValue(input.TableName)] {
		items = append(items, item)
	}
	return &dynamodb.QueryOutput{
		Items: items,
	}, nil
}

func TestGetInstanceList(t *testing.T) {
	mdb := NewMockDynamoDB()
	dao := &TaskDao{mdb}
	var err error

	instance0 := model.TaskInstance{
		ID:     "instance-001",
		Task:   "test-get-instance-task",
		Hook:   "http://example.com/instance-001",
		Status: 4,
		Tags: map[string]interface{}{
			"source": "auto",
		},
		Args: map[string]interface{}{
			"arg0": "the 1st arg",
			"arg1": float64(1),
		},
		Data: map[string]interface{}{
			"data0": "hello-ce",
		},
	}
	err = dao.PutInstance(&instance0)
	assert.Nil(t, err)

	instance1 := model.TaskInstance{
		ID:     "instance-002",
		Task:   "test-get-instance-task",
		Hook:   "http://example.com/instance-002",
		Status: 4,
		Tags: map[string]interface{}{
			"source": "auto",
		},
		Args: map[string]interface{}{
			"arg0": float64(0),
			"arg1": "the 2rd arg",
		},
		Data: map[string]interface{}{
			"data1": "hello world",
		},
	}
	err = dao.PutInstance(&instance1)
	assert.Nil(t, err)

	instannces, err := dao.GetInstanceList([]string{"instance-001", "instance-002"})
	assert.Nil(t, err)

	assert.Equal(t, 2, len(instannces))
	assert.NotEmpty(t, instannces[0].CreatedAt)
	assert.NotEmpty(t, instannces[0].UpdatedAt)
	assert.NotEmpty(t, instannces[1].CreatedAt)
	assert.NotEmpty(t, instannces[1].UpdatedAt)
	instannces[0].CreatedAt = model.JSONTime(time.Time{})
	instannces[0].UpdatedAt = model.JSONTime(time.Time{})
	instannces[1].CreatedAt = model.JSONTime(time.Time{})
	instannces[1].UpdatedAt = model.JSONTime(time.Time{})
	assert.Equal(t, instance0, instannces[0])
	assert.Equal(t, instance1, instannces[1])
}

func TestSearchInstance(t *testing.T) {
	mdb := NewMockDynamoDB()
	dao := &TaskDao{mdb}
	var err error

	instance0 := model.TaskInstance{
		ID:     "instance-001",
		Task:   "test-get-instance-task",
		Hook:   "http://example.com/instance-001",
		Status: 4,
		Tags: map[string]interface{}{
			"source": "auto",
		},
		Args: map[string]interface{}{
			"arg0": "the 1st arg",
			"arg1": float64(1),
		},
		Data: map[string]interface{}{
			"data0": "hello-ce",
		},
	}
	err = dao.PutInstance(&instance0)
	assert.Nil(t, err)

	instance1 := model.TaskInstance{
		ID:     "instance-002",
		Task:   "test-get-instance-task",
		Hook:   "http://example.com/instance-002",
		Status: 4,
		Tags: map[string]interface{}{
			"source": "auto",
		},
		Args: map[string]interface{}{
			"arg0": float64(0),
			"arg1": "the 2rd arg",
		},
		Data: map[string]interface{}{
			"data1": "hello world",
		},
	}
	err = dao.PutInstance(&instance1)
	assert.Nil(t, err)

	ids, err := dao.SearchInstance("test-get-instance-task", 4, nil, nil, nil, nil, nil)
	assert.Nil(t, err)
	assert.ElementsMatch(t, []string{"instance-001", "instance-002"}, ids)
}

func TestPutInstance(t *testing.T) {
	mdb := NewMockDynamoDB()
	dao := &TaskDao{mdb}
	var err error

	instance0 := model.TaskInstance{
		ID:     "instance-001",
		Task:   "test-get-instance-task",
		Hook:   "http://example.com/instance-001",
		Status: 4,
		Tags: map[string]interface{}{
			"source": "auto",
		},
		Args: map[string]interface{}{
			"arg0": "the 1st arg",
			"arg1": float64(1),
		},
		Data: map[string]interface{}{
			"data0": "hello-ce",
		},
	}
	err = dao.PutInstance(&instance0)
	assert.Nil(t, err)

	instance1 := model.TaskInstance{
		ID:     "instance-002",
		Task:   "test-get-instance-task",
		Hook:   "http://example.com/instance-002",
		Status: 4,
		Tags: map[string]interface{}{
			"source": "auto",
		},
		Args: map[string]interface{}{
			"arg0": float64(0),
			"arg1": "the 2rd arg",
		},
		Data: map[string]interface{}{
			"data1": "hello world",
		},
	}
	err = dao.PutInstance(&instance1)
	assert.Nil(t, err)

	instance1.Status = 8
	instance1.Data = map[string]interface{}{
		"error": "failed",
	}
	err = dao.PutInstance(&instance1)
	assert.Nil(t, err)

	assert.Equal(t, model.Failed, instance1.Status)
	assert.Equal(t, map[string]interface{}{
		"error": "failed",
	}, instance1.Data)

	instannces, err := dao.GetInstanceList([]string{"instance-001", "instance-002"})
	assert.Nil(t, err)

	instannces[0].CreatedAt = model.JSONTime(time.Time{})
	instannces[0].UpdatedAt = model.JSONTime(time.Time{})
	instannces[1].CreatedAt = model.JSONTime(time.Time{})
	instannces[1].UpdatedAt = model.JSONTime(time.Time{})
	assert.Equal(t, instance0, instannces[0])
	assert.Equal(t, instance1, instannces[1])
}