// Copyright 2022 Cisco Systems, Inc. and its affiliates
//
// 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,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package mongodb

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.uber.org/zap"

	"github.com/cisco-open/flame/cmd/controller/app/objects"
	"github.com/cisco-open/flame/pkg/openapi"
	"github.com/cisco-open/flame/pkg/util"
)

// CreateTasks creates task records in task db collection
func (db *MongoService) CreateTasks(tasks []objects.Task, dirty bool) error {
	zap.S().Debugf("Calling CreateTasks")

	success := false

	// rollback closure in case of error
	defer func() {
		if success {
			return
		}
		// TODO: implement this
	}()

	var updates []mongo.WriteModel

	for _, task := range tasks {
		cfgData, err := json.Marshal(&task.JobConfig)
		if err != nil {
			return fmt.Errorf("failed to marshal task: %v", err)
		}

		zap.S().Debugf("Creating task for task %s", task.TaskId)
		filter := bson.M{util.DBFieldJobId: task.JobId, util.DBFieldTaskId: task.TaskId}
		update := bson.M{
			"$set": bson.M{
				util.DBFieldRole:             task.Role,
				util.DBFieldGroupAssociation: task.JobConfig.GroupAssociation,
				"config":                     cfgData,
				"code":                       task.ZippedCode,
				util.DBFieldComputeId:        task.ComputeId,
				util.DBFieldTaskDirty:        dirty,
				util.DBFieldTaskKey:          task.Key,
				util.DBFieldState:            openapi.READY,
				util.DBFieldTimestamp:        time.Now(),
			},
		}
		zap.S().Debugf("taskID %s is assigned compute %s", task.TaskId, task.ComputeId)

		updateModel := mongo.NewUpdateOneModel().SetFilter(filter).SetUpdate(update).SetUpsert(true)
		updates = append(updates, updateModel)
	}

	opts := options.BulkWriteOptions{}
	_, err := db.taskCollection.BulkWrite(context.TODO(), updates, &opts)
	if err != nil {
		zap.S().Errorf("Failed to insert tasks: %v", err)
		return err
	}

	success = true
	zap.S().Debugf("Created tasks successufully")

	return nil
}

func (db *MongoService) GetTask(jobId string, taskId string, key string) (map[string][]byte, error) {
	zap.S().Infof("Fetching task - jobId: %s, taskId: %s", jobId, taskId)

	type taskInMongo struct {
		Config []byte `json:"config"`
		Code   []byte `json:"code"`
		Key    string `json:"key"`
	}
	var task taskInMongo

	filter := bson.M{util.DBFieldJobId: jobId, util.DBFieldTaskId: taskId}

	err := db.taskCollection.FindOne(context.TODO(), filter).Decode(&task)
	if err != nil {
		zap.S().Warnf("Failed to fetch task: %v", err)
		return nil, err
	}

	// The key in the system is not empty and it doesn't match with one provided
	if task.Key != "" && task.Key != key {
		err = fmt.Errorf("keys don't match")
		zap.S().Warnf("%v", err)
		return nil, err
	}

	taskMap := map[string][]byte{
		util.TaskConfigFile: task.Config,
		util.TaskCodeFile:   task.Code,
	}

	// This is the first time that task is registered; let's set the key in DB
	if task.Key == "" {
		setElements := bson.M{util.DBFieldTaskKey: key}

		update := bson.M{"$set": setElements}

		after := options.After
		upsert := true
		opts := options.FindOneAndUpdateOptions{
			ReturnDocument: &after,
			Upsert:         &upsert,
		}

		updatedDoc := openapi.TaskStatus{}
		err = db.taskCollection.FindOneAndUpdate(context.TODO(), filter, update, &opts).Decode(&updatedDoc)
		if err != nil {
			zap.S().Warnf("Failed to set key: %v", err)
			return nil, fmt.Errorf("failed to set key - %v", ErrorCheck(err))
		}
	}

	return taskMap, nil
}

func (db *MongoService) DeleteTasks(jobId string, dirty bool) error {
	zap.S().Infof("Deleting tasks for job: %s", jobId)
	filter := bson.M{util.DBFieldJobId: jobId, util.DBFieldTaskDirty: dirty}

	_, err := db.taskCollection.DeleteMany(context.TODO(), filter)
	if err != nil {
		zap.S().Warnf("Failed to remove task: %v", err)
		return err
	}

	return nil
}

func (db *MongoService) UpdateTaskStatus(jobId string, taskId string, taskStatus openapi.TaskStatus) error {
	switch taskStatus.State {
	case openapi.RUNNING:
		fallthrough
	case openapi.FAILED:
		fallthrough
	case openapi.TERMINATED:
		fallthrough
	case openapi.COMPLETED:

	case openapi.READY:
		fallthrough
	case openapi.STARTING:
		fallthrough
	case openapi.APPLYING:
		fallthrough
	case openapi.DEPLOYING:
		fallthrough
	case openapi.STOPPING:
		return fmt.Errorf("prohibited state: %s", taskStatus.State)

	default:
		return fmt.Errorf("unknown state: %s", taskStatus.State)
	}

	filter := bson.M{util.DBFieldJobId: jobId, util.DBFieldTaskId: taskId}

	setElements := bson.M{
		util.DBFieldState:     taskStatus.State,
		util.DBFieldTaskLog:   taskStatus.Log,
		util.DBFieldTimestamp: time.Now(),
	}

	update := bson.M{"$set": setElements}

	after := options.After
	upsert := true
	opts := options.FindOneAndUpdateOptions{
		ReturnDocument: &after,
		Upsert:         &upsert,
	}

	updatedDoc := openapi.TaskStatus{}
	err := db.taskCollection.FindOneAndUpdate(context.TODO(), filter, update, &opts).Decode(&updatedDoc)
	if err != nil {
		return ErrorCheck(err)
	}

	return nil
}

// UpdateTaskStateByFilter updates the status of tasks on a particular state in a job.
// The function should be only called by a system; hence, this function shouldn't become open to user API
func (db *MongoService) UpdateTaskStateByFilter(jobId string, newState openapi.JobState, userFilter map[string]interface{}) error {
	filter := bson.M{util.DBFieldJobId: jobId}
	for k, v := range userFilter {
		filter[k] = v
	}

	setElements := bson.M{util.DBFieldState: newState, util.DBFieldTimestamp: time.Now()}

	update := bson.M{"$set": setElements}

	_, err := db.taskCollection.UpdateMany(context.TODO(), filter, update)
	if err != nil {
		zap.S().Warnf("Failed to update task state with filter %v: %v", userFilter, err)
		return err
	}

	return nil
}

// IsOneTaskInState returns true if the state of at least one task is set to the given state.
// Otherwise, it returns false.
func (db *MongoService) IsOneTaskInState(jobId string, state openapi.JobState) bool {
	filter := bson.M{util.DBFieldJobId: jobId, util.DBFieldState: state}

	result := db.taskCollection.FindOne(context.TODO(), filter)

	if result.Err() != nil {
		return false
	}

	return true
}

func (db *MongoService) IsOneTaskInStateWithRole(jobId string, state openapi.JobState, role string) bool {
	filter := bson.M{util.DBFieldJobId: jobId, util.DBFieldState: state, util.DBFieldRole: role}

	result := db.taskCollection.FindOne(context.TODO(), filter)

	if result.Err() != nil {
		return false
	}

	return true
}

func (db *MongoService) SetTaskDirtyFlag(jobId string, dirty bool) error {
	zap.S().Infof("Setting dirty flag %t for tasks of job: %s", dirty, jobId)

	filter := bson.M{util.DBFieldJobId: jobId}
	update := bson.M{
		"$set": bson.M{
			util.DBFieldTaskDirty: dirty,
		},
	}

	_, err := db.taskCollection.UpdateMany(context.TODO(), filter, update)
	if err != nil {
		zap.S().Warnf("Failed to set dirty flag to %s: %v", dirty, err)
		return err
	}

	return nil
}

// MonitorTasks creates and returns a watcher to monitor state changes of tasks in a job
func (db *MongoService) MonitorTasks(jobId string) (chan openapi.TaskInfo, chan error, context.CancelFunc, error) {
	jobIdField := fmt.Sprintf("fullDocument.%s", util.DBFieldJobId)
	stateField := fmt.Sprintf("updateDescription.updatedFields.%s", util.DBFieldState)

	pipeline := mongo.Pipeline{
		bson.D{{
			Key: "$match",
			Value: bson.D{{
				Key: "$and",
				Value: bson.A{
					bson.D{{Key: jobIdField, Value: jobId}},
					bson.D{{Key: stateField, Value: bson.M{"$exists": true}}},
					bson.D{{Key: "operationType", Value: "update"}},
				},
			}},
		}},
	}

	chanLen := 100
	opts := options.ChangeStream().SetFullDocument(options.UpdateLookup)
	watcher := db.newstreamWatcher(db.taskCollection, pipeline, opts, chanLen)
	if err := watcher.watch(); err != nil {
		return nil, nil, nil, err
	}

	eventCh := make(chan openapi.TaskInfo, chanLen)

	go func() {
		isDone := false

		for !isDone {
			select {
			case event := <-watcher.streamCh:
				fullDoc := event["fullDocument"].(bson.M)
				data, err := json.Marshal(fullDoc)
				if err != nil {
					zap.S().Debugf("Failed to marshal event: %v", err)
					watcher.errCh <- fmt.Errorf("marshaling failed")
					isDone = true
					continue
				}

				var tskEvent openapi.TaskInfo
				err = json.Unmarshal(data, &tskEvent)
				if err != nil {
					zap.S().Debugf("Failed to unmarshal to task event: %v", err)
					watcher.errCh <- fmt.Errorf("unmarshaling failed")
					isDone = true
					continue
				}

				eventCh <- tskEvent

			case <-watcher.ctx.Done():
				isDone = true
			}
		}

		zap.S().Infof("Finished monitoring tasks for job %s", jobId)
	}()

	return eventCh, watcher.errCh, watcher.cancel, nil
}
