package impl

import (
	"bytes"
	"context"
	"encoding/csv"
	"io"
	"testing"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	s3types "github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/stretchr/testify/require"
)

func TestCountingWriterTracksManifestBytes(t *testing.T) {
	var output bytes.Buffer
	countingOutput := &countingWriter{Writer: &output}
	writer := csv.NewWriter(countingOutput)

	require.NoError(t, writer.Write([]string{"source-bucket", "prefix/object,with-comma"}))
	writer.Flush()
	require.NoError(t, writer.Error())

	require.Equal(t, int64(output.Len()), countingOutput.bytes)
	require.Equal(t, "source-bucket,\"prefix/object,with-comma\"\n", output.String())
}

func TestWriteManifestPagesStreamsLargeObjectList(t *testing.T) {
	const objectCount = 100_000
	objects := make([]s3types.Object, objectCount+1)
	objects[0] = s3types.Object{Key: aws.String("prefix/")}
	for i := 1; i <= objectCount; i++ {
		objects[i] = s3types.Object{Key: aws.String("prefix/object")}
	}
	paginator := &fakeManifestPaginator{pages: []*s3.ListObjectsV2Output{{Contents: objects}}}

	stats, err := writeManifestPages(context.Background(), io.Discard, "source-bucket", paginator)

	require.NoError(t, err)
	require.Equal(t, int64(objectCount), stats.objects)
	require.Equal(t, int64(objectCount*len("source-bucket,prefix/object\n")), stats.bytes)
}

type fakeManifestPaginator struct {
	pages []*s3.ListObjectsV2Output
	index int
}

func (p *fakeManifestPaginator) HasMorePages() bool {
	return p.index < len(p.pages)
}

func (p *fakeManifestPaginator) NextPage(context.Context) (*s3.ListObjectsV2Output, error) {
	page := p.pages[p.index]
	p.index++
	return page, nil
}