package s3

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"sync"

	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	s3api "github.com/aws/aws-sdk-go-v2/service/s3"
	s3type "github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/pingcap/log"
	"github.com/tidbcloud/observability/pkg/common/constants"
	"github.com/tidbcloud/observability/pkg/common/errors"
	"github.com/tidbcloud/observability/pkg/common/structs"
	"github.com/tidbcloud/observability/pkg/metamanager"
	"go.uber.org/zap"
	"gopkg.in/yaml.v3"
)

type modelMeta struct {
	awsClient *s3api.Client
}

var (
	modelMetaSvc     metamanager.ModelMeta
	modelMetaSvcOnce sync.Once
)

func GetModelMetaService() metamanager.ModelMeta {
	modelMetaSvcOnce.Do(func() {
		if modelMetaSvc == nil {
			cfg, err := awsconfig.LoadDefaultConfig(context.TODO())
			if err != nil {
				panic(err)
			}
			cfg.Region = "us-west-2"
			modelMetaSvc = &modelMeta{
				awsClient: s3api.NewFromConfig(cfg),
			}
		}
	})

	return modelMetaSvc
}

func (e modelMeta) SaveMeta(ctx context.Context, meta structs.Model) (err error) {
	bucket := constants.EvaluationBucketName()
	fileKey := fmt.Sprintf(constants.ModelMetaFileKey, meta.Name, meta.Version)

	metaYaml, err := yaml.Marshal(meta)
	if err != nil {
		return err
	}

	reader := bytes.NewReader(metaYaml)
	_, err = e.awsClient.PutObject(ctx, &s3api.PutObjectInput{
		Bucket: &bucket,
		Key:    &fileKey,
		Body:   reader,
	})
	if err != nil {
		err = errors.Wrap(errors.ModelWriteError, "put model meta to s3 failed", err)
		log.Error("put model meta to s3 failed", zap.Error(err))
		return
	}

	return
}

func (e modelMeta) GetMeta(ctx context.Context, name, version string) (model structs.Model, err error) {
	bucket := constants.EvaluationBucketName()
	fileKey := fmt.Sprintf(constants.ModelMetaFileKey, name, version)

	output, err := e.awsClient.GetObject(ctx, &s3api.GetObjectInput{
		Bucket: &bucket,
		Key:    &fileKey,
	})
	if err != nil {
		log.Error("get model meta from s3 failed", zap.Error(err))

		switch err.(type) {
		case *s3type.NoSuchBucket, *s3type.NoSuchKey, *s3type.NotFound:
			err = errors.Wrap(errors.ModelNotFound, "model meta not found from s3", err)
			return
		}

		err = errors.Wrap(errors.ModelReadError, "get model meta from s3 failed", err)
		return
	}

	defer output.Body.Close()

	content, err := io.ReadAll(output.Body)
	if err != nil {
		err = errors.Wrap(errors.ModelReadError, "read model meta from s3 failed", err)
		log.Error("read model meta from s3 failed", zap.Error(err))
		return
	}

	err = yaml.Unmarshal(content, &model)
	if err != nil {
		err = errors.Wrap(errors.ModelReadError, "parse model meta yaml failed", err)
		log.Error("parse model meta yaml failed", zap.Error(err))
		return
	}
	return
}

func MockModelMeta(mock metamanager.ModelMeta) {
	modelMetaSvc = mock
}