package sinks

import (
	"bytes"
	"net/http"
	"os"
	"testing"

	"github.com/aws/aws-sdk-go-v2/aws"
	awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	smithyhttp "github.com/aws/smithy-go/transport/http"
	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"
	"github.com/tidbcloud/observability/data-collector/pkg/collector/config"
	mock_s3 "github.com/tidbcloud/observability/data-collector/test/mocks3"
)

func newMockS3Sink(cfg config.S3Sink, awscfg aws.Config, mock S3Client) S3Sink {
	return S3Sink{
		cfg:    cfg,
		awscfg: awscfg,
		s3cli:  mock,
	}
}
func TestWriteToFile(t *testing.T) {
	ctl := gomock.NewController(t)
	var awscfg aws.Config
	awscfg.Region = "us-west-2"
	var cfg config.S3Sink
	mockclient := mock_s3.NewMockS3Client(ctl)
	s3sink := newMockS3Sink(cfg, awscfg, mockclient)
	mockclient.EXPECT().PutObject(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
	err := s3sink.WriteToFile("test", []byte("test"))
	assert.NoError(t, err)
}

func TestCheckFileExits(t *testing.T) {
	ctl := gomock.NewController(t)
	var awscfg aws.Config
	awscfg.Region = "us-west-2"
	var cfg config.S3Sink
	mockclient := mock_s3.NewMockS3Client(ctl)
	s3sink := newMockS3Sink(cfg, awscfg, mockclient)
	//success and file exists
	mockclient.EXPECT().HeadObject(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
	_, _, err := s3sink.CheckFileExits("test")
	assert.NoError(t, err)
	//file not found
	responseError := &awshttp.ResponseError{
		ResponseError: &smithyhttp.ResponseError{
			Response: &smithyhttp.Response{
				Response: &http.Response{
					StatusCode: 404,
				},
			},
		},
	}
	mockclient.EXPECT().HeadObject(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, responseError)
	_, exits, err := s3sink.CheckFileExits("test")
	assert.NoError(t, err)
	assert.False(t, exits)
	//other error
	mockclient.EXPECT().HeadObject(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, os.ErrInvalid)
	_, _, err = s3sink.CheckFileExits("test")
	assert.Error(t, err)

}

type CBuffer struct {
	*bytes.Buffer
}

func (c *CBuffer) Close() error {
	return nil
}

func TestDownloadFile(t *testing.T) {
	ctl := gomock.NewController(t)
	var awscfg aws.Config
	awscfg.Region = "us-west-2"
	var cfg config.S3Sink
	mockclient := mock_s3.NewMockS3Client(ctl)
	s3sink := newMockS3Sink(cfg, awscfg, mockclient)
	mockclient.EXPECT().HeadObject(gomock.Any(), gomock.Any(), gomock.Any()).Return(&s3.HeadObjectOutput{
		ContentLength: 1,
	}, nil)
	buf := make([]byte, 1)
	tmp := CBuffer{bytes.NewBuffer(buf)}
	mockclient.EXPECT().GetObject(gomock.Any(), gomock.Any(), gomock.Any()).Return(&s3.GetObjectOutput{
		Body:          &tmp,
		ContentLength: 1,
	}, nil)
	_, _, err := s3sink.DownloadFile("test")
	assert.NoError(t, err)
	//error
	mockclient.EXPECT().HeadObject(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, os.ErrInvalid)
	_, _, err = s3sink.DownloadFile("test")
	assert.Error(t, err)
	//download error
	mockclient.EXPECT().HeadObject(gomock.Any(), gomock.Any(), gomock.Any()).Return(&s3.HeadObjectOutput{
		ContentLength: 1,
	}, nil)
	mockclient.EXPECT().GetObject(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, os.ErrInvalid)
	_, _, err = s3sink.DownloadFile("test")
	assert.Error(t, err)
}

func TestNewSink(t *testing.T) {
	os.Setenv("AWS_ACCESS_KEY_ID", "test1")
	os.Setenv("AWS_SECRET_ACCESS_KEY", "testsec")
	defer os.Unsetenv("AWS_SECRET_ACCESS_KEY")
	defer os.Unsetenv("AWS_ACCESS_KEY_ID")
	var cfg config.S3Sink
	assert.NotNil(t, NewS3Sink(cfg))
}
