package storage

import (
	"bytes"
	"context"
	"errors"
	"io"
	"testing"
)

func TestReadObjectBoundedAcceptsExactLimitWithAbsentSizeMetadata(t *testing.T) {
	body := &trackingReadCloser{reader: bytes.NewReader([]byte("12345678")), chunkSize: 2}
	store := &boundedReadStore{read: func(context.Context, string) (ObjectInfo, io.ReadCloser, error) {
		return ObjectInfo{Key: "control.json", ETag: "generation"}, body, nil
	}}

	info, data, err := ReadObjectBounded(context.Background(), store, "control.json", 8)
	if err != nil {
		t.Fatalf("ReadObjectBounded() error = %v", err)
	}
	if string(data) != "12345678" || info.ETag != "generation" {
		t.Fatalf("ReadObjectBounded() = (%#v, %q), want opened object metadata and bytes", info, data)
	}
	if !body.closed {
		t.Fatal("object body was not closed")
	}
}

func TestReadObjectBoundedRejectsLimitPlusOneDespiteMissingOrInaccurateMetadata(t *testing.T) {
	tests := []struct {
		name         string
		metadataSize int64
	}{
		{name: "missing", metadataSize: 0},
		{name: "inaccurate", metadataSize: 2},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			body := &trackingReadCloser{reader: bytes.NewReader([]byte("123456789")), chunkSize: 1}
			store := &boundedReadStore{read: func(context.Context, string) (ObjectInfo, io.ReadCloser, error) {
				return ObjectInfo{Key: "control.json", Size: test.metadataSize}, body, nil
			}}

			_, data, err := ReadObjectBounded(context.Background(), store, "control.json", 8)
			var limitErr *ReadLimitError
			if !errors.As(err, &limitErr) {
				t.Fatalf("ReadObjectBounded() error = %v, want ReadLimitError", err)
			}
			if data != nil || body.bytesRead != 9 || !body.closed {
				t.Fatalf("data=%q bytes read=%d closed=%t, want nil, 9, true", data, body.bytesRead, body.closed)
			}
		})
	}
}

func TestReadObjectBoundedRejectsOversizedMetadataBeforeTransfer(t *testing.T) {
	body := &trackingReadCloser{reader: bytes.NewReader([]byte("small"))}
	store := &boundedReadStore{read: func(context.Context, string) (ObjectInfo, io.ReadCloser, error) {
		return ObjectInfo{Key: "control.json", Size: 9}, body, nil
	}}

	_, _, err := ReadObjectBounded(context.Background(), store, "control.json", 8)
	var limitErr *ReadLimitError
	if !errors.As(err, &limitErr) {
		t.Fatalf("ReadObjectBounded() error = %v, want ReadLimitError", err)
	}
	if body.bytesRead != 0 || !body.closed {
		t.Fatalf("bytes read=%d closed=%t, want zero-byte transfer and closed body", body.bytesRead, body.closed)
	}
}

func TestReadObjectBoundedPropagatesCancellationAndClosesBody(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	body := &trackingReadCloser{reader: bytes.NewReader([]byte("12345678")), chunkSize: 1, afterRead: cancel}
	store := &boundedReadStore{read: func(context.Context, string) (ObjectInfo, io.ReadCloser, error) {
		return ObjectInfo{Key: "control.json"}, body, nil
	}}

	_, data, err := ReadObjectBounded(ctx, store, "control.json", 8)
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("ReadObjectBounded() error = %v, want context cancellation", err)
	}
	if data != nil || body.bytesRead != 1 || !body.closed {
		t.Fatalf("data=%q bytes read=%d closed=%t, want nil, 1, true", data, body.bytesRead, body.closed)
	}
}

func TestReadObjectBoundedReturnsCloseFailure(t *testing.T) {
	closeErr := errors.New("close failed")
	body := &trackingReadCloser{reader: bytes.NewReader([]byte("ok")), closeErr: closeErr}
	store := &boundedReadStore{read: func(context.Context, string) (ObjectInfo, io.ReadCloser, error) {
		return ObjectInfo{Key: "control.json", Size: 2}, body, nil
	}}

	_, data, err := ReadObjectBounded(context.Background(), store, "control.json", 8)
	if !errors.Is(err, closeErr) || data != nil || !body.closed {
		t.Fatalf("data=%q error=%v closed=%t, want close failure and no retained data", data, err, body.closed)
	}
}

type boundedReadStore struct {
	ObjectStore
	read func(context.Context, string) (ObjectInfo, io.ReadCloser, error)
}

func (s *boundedReadStore) Read(ctx context.Context, key string) (ObjectInfo, io.ReadCloser, error) {
	return s.read(ctx, key)
}

type trackingReadCloser struct {
	reader    io.Reader
	chunkSize int
	afterRead func()
	closeErr  error
	bytesRead int
	closed    bool
}

func (r *trackingReadCloser) Read(p []byte) (int, error) {
	if r.chunkSize > 0 && len(p) > r.chunkSize {
		p = p[:r.chunkSize]
	}
	n, err := r.reader.Read(p)
	r.bytesRead += n
	if n > 0 && r.afterRead != nil {
		r.afterRead()
		r.afterRead = nil
	}
	return n, err
}

func (r *trackingReadCloser) Close() error {
	r.closed = true
	return r.closeErr
}