package decrypt

import (
	"bytes"
	"crypto/rsa"
	"errors"
	"fmt"
	"io"

	"github.com/pingcap-inc/clinic/pkg/auth/session"
	"github.com/pingcap-inc/clinic/pkg/datasource"
	"github.com/pingcap-inc/clinic/pkg/resource_limit"

	"github.com/pingcap-inc/clinic/pkg/api/common"
	"github.com/pingcap-inc/clinic/pkg/storage/aws"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/s3"

	"go.uber.org/zap"

	"github.com/pingcap-inc/clinic/pkg/config"
	_const "github.com/pingcap-inc/clinic/pkg/const"
	"github.com/pingcap-inc/clinic/pkg/es"
	"github.com/pingcap-inc/clinic/pkg/executor"
	"github.com/pingcap-inc/clinic/pkg/metrics"

	"github.com/pingcap-inc/clinic/pkg/storage/iface"

	"github.com/pingcap/diag/pkg/crypto"
	"github.com/pingcap/log"

	"github.com/pingcap-inc/clinic/pkg/db/model"
	task_manager "github.com/pingcap-inc/clinic/pkg/task_manager"
)

var decryptExecutor = executor.NewExecutor(10)

type Decrypt struct {
	privRSAKey     *rsa.PrivateKey
	fileStorage    iface.IStorage
	TaskManagerCli *task_manager.TaskManagerClient
	quotaLimiter   *resource_limit.QuotaLimiter
}

func NewDecrypt(privRSAKey *rsa.PrivateKey, fileStorage iface.IStorage, taskManagerCli *task_manager.TaskManagerClient, quotaLimiter *resource_limit.QuotaLimiter) *Decrypt {
	return &Decrypt{
		privRSAKey:     privRSAKey,
		fileStorage:    fileStorage,
		TaskManagerCli: taskManagerCli,
		quotaLimiter:   quotaLimiter,
	}
}

func mayRetry(err error) bool {
	if errors.Is(err, rsa.ErrDecryption) {
		return false
	}
	if errors.Is(err, rsa.ErrVerification) {
		return false
	}
	// may be read buffer failed error, can retry.
	return true
}

func (d *Decrypt) decryptCachedFile(fileUUID, clusterID string) error {
	buildDecryptor := func(fileUUID string) (crypto.Decryptor, io.Closer, error) {
		resp, err := d.fileStorage.GetObject(fileUUID, aws.TmpPrefix, iface.GetObjectRequest{})
		if err != nil {
			return nil, nil, err
		}

		dec, err := crypto.NewDecryptor(d.privRSAKey, resp.Body)
		if err != nil {
			// close the resp body on error.
			// NewDecryptor already consume the body, no need to drain body again
			resp.Body.Close()
			return nil, nil, err
		}
		return dec, resp.Body, nil
	}
	var (
		dec    crypto.Decryptor
		err    error
		closer io.Closer
	)
	dec, closer, err = buildDecryptor(fileUUID)
	if err != nil {
		// here we need some retry logic to handle decrypt error caused by broken buffer.
		if !mayRetry(err) {
			return err
		}
		if dec, closer, err = buildDecryptor(fileUUID); err != nil {
			return err
		}
	}
	defer func() {
		if err := closer.Close(); err != nil {
			log.Error("decryptor closer close failed", zap.Error(err))
		}
	}()

	err = d.fileStorage.Persist(fileUUID, clusterID, dec)
	if err != nil {
		return err
	}

	return d.fileStorage.DeleteCache(fileUUID)
}

type decryptTask struct {
	decrypt   *Decrypt
	fileUUID  string
	clusterID string
	rebuild   bool
}

func (dt *decryptTask) Run() {
	metrics.WaitDecryptionTaskCount.Dec()
	err := dt.decrypt.decryptCachedFile(dt.fileUUID, dt.clusterID)
	if err != nil {
		log.Error("decode file failed", zap.String("file", dt.fileUUID), zap.Error(err))
		return
	}

	// default rebuild: metrics and conf
	if dt.rebuild {
		if err := dt.decrypt.TaskManagerCli.Rebuild(common.Rebuild{DataType: int(_const.MetricsAndConfTask), ItemID: dt.fileUUID, HealthReport: dt.rebuild}); err != nil {
			log.Error("start rebuild failed")
		}
	}

	log.Info("decode file success", zap.String("file", dt.fileUUID))
	// TODO: clean oldest file, the error should not broken the whole process
	//if err := dt.cleanOldestFile(); err != nil {
	//	log.Error("clean oldest file failed", zap.Error(err))
	//}
}

func (d *Decrypt) DecryptString(ciphertext []byte) ([]byte, error) {
	dec, err := crypto.NewDecryptor(d.privRSAKey, bytes.NewBuffer(ciphertext))
	if err != nil {
		return nil, err
	}
	plaintext := bytes.NewBuffer([]byte{})
	if _, err := io.Copy(plaintext, dec); err != nil {
		return nil, err
	}
	return plaintext.Bytes(), nil
}

func (dt *decryptTask) cleanOldestFile() error {
	uf, err := model.QueryUploadFile(dt.fileUUID)
	if err != nil {
		log.Error("query upload file failed", zap.String("file", dt.fileUUID), zap.Error(err))
		return err
	}

	// skip internal user
	org, err := datasource.GetOrganization(&session.CtxUser{
		Internal: true,
		Email:    "clinic@pingcap.com",
	}, uf.ID)
	if err != nil {
		log.Error("query org failed", zap.String("file", dt.fileUUID), zap.Error(err))
		return err
	}
	user, err := model.GetUserByID(org.CreatedBy)
	if err != nil {
		log.Error("query user failed", zap.String("file", dt.fileUUID), zap.Error(err))
		return err
	}
	if user.Internal {
		return nil
	}

	uploadedBytes, err := model.FetchTotalClusterDataSize(uf.ClusterID)
	if err != nil {
		log.Error("cal file total bytes failed", zap.Error(err))
		return err
	}
	quota := uint64(_const.MaxFileSizePreClusterInBytes)
	rl, err := dt.decrypt.quotaLimiter.GetResourceLimiter(resource_limit.DataSizeQuota)
	if err != nil {
		log.Error("get resource limiter failed", zap.Stringer("resource", resource_limit.DataSizeQuota), zap.Error(err))
	} else {
		quota = rl.GetQuota(uf.OrgID)
	}
	if uploadedBytes <= quota {
		return nil
	}
	oldestFile, found, err := model.FindOldestCompletedFileByClusterID(uf.ClusterID)
	if err != nil {
		log.Error("find oldest file failed", zap.Error(err))
		return err
	}
	if !found {
		return nil
	}
	if oldestFile.FileUUID == uf.FileUUID {
		return nil
	}
	var deleteResourceErr error
	// TODO: can use goroutines to make the process parallel
	if err := dt.decrypt.fileStorage.DeleteObject(oldestFile.FileUUID, oldestFile.ClusterID); err != nil {
		if aerr, ok := err.(awserr.Error); ok && aerr.Code() != s3.ErrCodeNoSuchKey {
			log.Error("delete object failed", zap.String("file", oldestFile.FileUUID), zap.Error(err))
			deleteResourceErr = err
		}
	}
	logIndex := fmt.Sprintf("org%d_log", uf.OrgID)
	if err := es.DeleteByFileUUID(logIndex, oldestFile.FileUUID, config.GetConfig().SinkConfig.LogConfig); err != nil {
		log.Error("delete log failed", zap.String("file", oldestFile.FileUUID), zap.Error(err))
		deleteResourceErr = err
	}
	slowlogIndex := fmt.Sprintf("org%d_slowlog", uf.OrgID)
	if err := es.DeleteByFileUUID(slowlogIndex, oldestFile.FileUUID, config.GetConfig().SinkConfig.LogConfig); err != nil {
		log.Error("delete slow log failed", zap.String("file", oldestFile.FileUUID), zap.Error(err))
		deleteResourceErr = err
	}
	// we should not delete data form vm
	// file record should not be deleted if any previous deletion failed
	if deleteResourceErr != nil {
		return deleteResourceErr
	}
	if err := model.DeleteFileRecord(oldestFile.FileUUID); err != nil {
		log.Error("delete file record failed", zap.String("file", oldestFile.FileUUID), zap.Error(err))
	}
	return nil
}

func (dt *Decrypt) SubmitTask(fileUUID, clusterID string, rebuild bool) {
	metrics.WaitDecryptionTaskCount.Inc()
	decryptExecutor.Execute(&decryptTask{
		decrypt:   dt,
		fileUUID:  fileUUID,
		clusterID: clusterID,
		rebuild:   rebuild,
	})
}