package objectstore

import (
	"bytes"
	"context"
	"io"
	"testing"
	"time"

	"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"
)

type fakeS3V2Client struct {
	putInput   *s3.PutObjectInput
	headInput  *s3.HeadObjectInput
	getInput   *s3.GetObjectInput
	listInputs []*s3.ListObjectsV2Input
	body       []byte
}

func (f *fakeS3V2Client) PutObject(ctx context.Context, input *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
	f.putInput = input
	data, err := io.ReadAll(input.Body)
	if err != nil {
		return nil, err
	}
	f.body = data
	return &s3.PutObjectOutput{ETag: aws.String("etag-1")}, nil
}

func (f *fakeS3V2Client) GetObject(ctx context.Context, input *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error) {
	f.getInput = input
	return &s3.GetObjectOutput{
		Body:          io.NopCloser(bytes.NewReader(f.body)),
		ContentLength: aws.Int64(int64(len(f.body))),
		ETag:          aws.String("etag-1"),
		Metadata:      map[string]string{"m": "v"},
	}, nil
}

func (f *fakeS3V2Client) HeadObject(ctx context.Context, input *s3.HeadObjectInput, optFns ...func(*s3.Options)) (*s3.HeadObjectOutput, error) {
	f.headInput = input
	now := time.Date(2026, 6, 25, 0, 0, 0, 0, time.UTC)
	return &s3.HeadObjectOutput{
		ContentLength: aws.Int64(int64(len(f.body))),
		ETag:          aws.String("etag-1"),
		LastModified:  &now,
		Metadata:      map[string]string{"m": "v"},
	}, nil
}

func (f *fakeS3V2Client) ListObjectsV2(ctx context.Context, input *s3.ListObjectsV2Input, optFns ...func(*s3.Options)) (*s3.ListObjectsV2Output, error) {
	f.listInputs = append(f.listInputs, input)
	if input.ContinuationToken == nil {
		return &s3.ListObjectsV2Output{
			Contents:              []s3types.Object{{Key: aws.String("prefix/a"), Size: aws.Int64(10), ETag: aws.String("e")}},
			IsTruncated:           aws.Bool(true),
			NextContinuationToken: aws.String("next"),
		}, nil
	}
	return &s3.ListObjectsV2Output{
		Contents: []s3types.Object{{Key: aws.String("prefix/b"), Size: aws.Int64(20), ETag: aws.String("e2")}},
	}, nil
}

func TestS3V2StorePutUsesPrefixAndMetadata(t *testing.T) {
	client := &fakeS3V2Client{}
	store := NewS3V2Store(client, "bucket", "prefix/")

	info, err := store.Put(context.Background(), "/a/b.jsonl.gz", bytes.NewReader([]byte("data")), PutOptions{
		ContentType:     "application/x-ndjson",
		ContentEncoding: "gzip",
		Metadata:        map[string]string{"schema-version": "1"},
	})
	require.NoError(t, err)
	require.Equal(t, "prefix/a/b.jsonl.gz", info.Key)
	require.Equal(t, "bucket", aws.ToString(client.putInput.Bucket))
	require.Equal(t, "prefix/a/b.jsonl.gz", aws.ToString(client.putInput.Key))
	require.Equal(t, "application/x-ndjson", aws.ToString(client.putInput.ContentType))
	require.Equal(t, "gzip", aws.ToString(client.putInput.ContentEncoding))
	require.Equal(t, map[string]string{"schema-version": "1"}, client.putInput.Metadata)
	require.Equal(t, []byte("data"), client.body)
}

func TestS3V2StoreHeadAndGet(t *testing.T) {
	client := &fakeS3V2Client{body: []byte("data")}
	store := NewS3V2Store(client, "bucket", "prefix")

	info, ok, err := store.Head(context.Background(), "a")
	require.NoError(t, err)
	require.True(t, ok)
	require.Equal(t, int64(4), info.Size)
	require.Equal(t, "prefix/a", aws.ToString(client.headInput.Key))

	body, getInfo, err := store.Get(context.Background(), "a")
	require.NoError(t, err)
	defer body.Close()
	data, err := io.ReadAll(body)
	require.NoError(t, err)
	require.Equal(t, []byte("data"), data)
	require.Equal(t, "prefix/a", getInfo.Key)
	require.Equal(t, "prefix/a", aws.ToString(client.getInput.Key))

	body, _, err = store.Get(context.Background(), "prefix/a")
	require.NoError(t, err)
	body.Close()
	require.Equal(t, "prefix/a", aws.ToString(client.getInput.Key))
}

func TestS3V2StoreListUsesPagination(t *testing.T) {
	client := &fakeS3V2Client{}
	store := NewS3V2Store(client, "bucket", "prefix")

	objects, err := store.List(context.Background(), "a")
	require.NoError(t, err)
	require.Len(t, objects, 2)
	require.Equal(t, "prefix/a", objects[0].Key)
	require.Equal(t, "prefix/b", objects[1].Key)
	require.Len(t, client.listInputs, 2)
	require.Equal(t, "prefix/a", aws.ToString(client.listInputs[0].Prefix))
	require.Nil(t, client.listInputs[0].ContinuationToken)
	require.Equal(t, "next", aws.ToString(client.listInputs[1].ContinuationToken))
}