package sinks

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"net/http"

	"github.com/aws/aws-sdk-go-v2/aws"
	awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/tidbcloud/observability/data-collector/pkg/collector/config"
)

type S3Client interface {
	PutObject(context.Context, *s3.PutObjectInput, ...func(*s3.Options)) (*s3.PutObjectOutput, error)
	HeadObject(context.Context, *s3.HeadObjectInput, ...func(*s3.Options)) (*s3.HeadObjectOutput, error)
	GetObject(context.Context, *s3.GetObjectInput, ...func(*s3.Options)) (*s3.GetObjectOutput, error)
}

type S3Sink struct {
	cfg    config.S3Sink
	awscfg aws.Config
	s3cli  S3Client
}

func NewS3Sink(cfg config.S3Sink) S3Sink {
	awscfg, err := awsconfig.LoadDefaultConfig(context.TODO(), awsconfig.WithRegion(cfg.Region))
	if err != nil {
		panic(fmt.Errorf("failed to load aws config, %w", err))
	}

	return S3Sink{
		cfg:    cfg,
		awscfg: awscfg,
		s3cli:  s3.NewFromConfig(awscfg),
	}
}

func (s3s *S3Sink) WriteToFile(path string, body []byte) error {
	rd := bytes.NewReader(body)
	_, err := s3s.s3cli.PutObject(context.TODO(), &s3.PutObjectInput{
		Bucket: aws.String(s3s.cfg.BucketName),
		Key:    aws.String(s3s.cfg.Prefix + path),
		Body:   rd,
	})
	return err
}

func (s3s *S3Sink) CheckFileExits(path string) (*s3.HeadObjectOutput, bool, error) {
	out, err := s3s.s3cli.HeadObject(context.TODO(), &s3.HeadObjectInput{
		Bucket: aws.String(s3s.cfg.BucketName),
		Key:    aws.String(s3s.cfg.Prefix + path),
	})
	if err != nil {
		var responseError *awshttp.ResponseError
		if errors.As(err, &responseError) && responseError.ResponseError.HTTPStatusCode() == http.StatusNotFound {
			return nil, false, nil
		}
		return nil, false, err
	}
	return out, true, nil
}

func (s3s *S3Sink) DownloadFile(path string) ([]byte, bool, error) {
	// Download the S3 object using the S3 manager object downloader
	downloader := manager.NewDownloader(s3s.s3cli)
	headObject, exits, err := s3s.CheckFileExits(path)
	if err != nil || !exits {
		return nil, exits, err
	}
	// pre-allocate in memory buffer, where headObject type is *s3.HeadObjectOutput
	buf := make([]byte, int(headObject.ContentLength))
	// wrap with aws.WriteAtBuffer
	w := manager.NewWriteAtBuffer(buf)
	_, err = downloader.Download(context.TODO(), w, &s3.GetObjectInput{
		Bucket: aws.String(s3s.cfg.BucketName),
		Key:    aws.String(s3s.cfg.Prefix + path),
	})
	if err != nil {
		return nil, false, err
	}
	return w.Bytes(), true, nil
}