package s3storage

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/aliyun/alibabacloud-oss-go-sdk-v2/oss"
	log "github.com/sirupsen/logrus"
	"gopkg.in/yaml.v3"
)

const checkpointFileName = "checkpoint.yaml"

type S3Checkpoint struct {
	bucket string
	prefix string
	cli    Client
}

type Data struct {
	Checkpoint int64 `json:"checkpoint" yaml:"checkpoint"`
	UpdatedAt  int64 `json:"updatedAt" yaml:"updatedAt"`
}

func BuildS3Checkpoint(bucket string, prefix string, cli Client) S3Checkpoint {
	return S3Checkpoint{
		bucket: bucket,
		prefix: prefix,
		cli:    cli,
	}
}

func (s S3Checkpoint) GetCheckpoint(ctx context.Context, taskName string) (int64, error) {
	s3FileName := fmt.Sprintf("%s/%s/%s", s.prefix, taskName, checkpointFileName)
	out, err := s.cli.Get(ctx, s.bucket, s3FileName)
	if err != nil {
		if strings.Contains(err.Error(), "StatusCode: 404") || strings.Contains(err.Error(), "object doesn't exist") {
			log.Infof("No checkpoint exists")
			return 0, nil
		}
		var serviceError *oss.ServiceError
		if errors.As(err, &serviceError) && (serviceError.Code == "NoSuchKey" || serviceError.StatusCode == http.StatusNotFound) {
			log.Infof("No checkpoint exists")
			return 0, nil
		}
		return 0, err
	}

	defer func() {
		if out != nil {
			(*out).Close()
		}
	}()

	content, err := io.ReadAll(*out)
	if err != nil {
		log.Errorf("Read checkpoint file failed, s3FileName = %s, err = %s", s3FileName, err.Error())
		return 0, err
	}

	data := Data{}
	err = yaml.Unmarshal(content, &data)
	if err != nil {
		log.Errorf("Unmarshal checkpoint file failed, content = %s, err = %s", string(content), err.Error())
		return 0, err
	}
	return data.Checkpoint, nil
}

func (s S3Checkpoint) SaveCheckpoint(ctx context.Context, taskName string, checkpointTimestamp int64) error {
	data := &Data{
		Checkpoint: checkpointTimestamp,
		UpdatedAt:  time.Now().Unix(),
	}

	content, err := yaml.Marshal(data)
	if err != nil {
		log.Errorf("Marshal checkpoint file failed, checkpointTimestamp = %d, err = %s", checkpointTimestamp, err.Error())
		return err
	}

	s3FileName := fmt.Sprintf("%s/%s/%s", s.prefix, taskName, checkpointFileName)

	err = s.cli.Put(ctx, s.bucket, s3FileName, bytes.NewReader(content))
	if err != nil {
		log.Errorf("Save checkpoint file failed, checkpointTimestamp = %d, err = %s", checkpointTimestamp, err.Error())
		return err
	}
	return nil
}