package utils

import (
	"fmt"
	"io"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
)

const (
	bufferSize = 4 * 1024 * 1024 // 4MB
)

type S3ReadSeeker struct {
	client     *s3.S3
	bucket     string
	key        string
	offset     int64
	body       io.ReadCloser
	contentLen int64
	buffer     []byte
	bufLen     int
	bufPos     int
}

func NewS3ReadSeeker(client *s3.S3, bucket, key string, body io.ReadCloser, contentLen int64) *S3ReadSeeker {
	return &S3ReadSeeker{
		client:     client,
		bucket:     bucket,
		key:        key,
		body:       body,
		contentLen: contentLen,
		buffer:     make([]byte, bufferSize),
	}
}

func (r *S3ReadSeeker) Read(p []byte) (n int, err error) {
	if r.bufPos < r.bufLen {
		n = copy(p, r.buffer[r.bufPos:r.bufLen])
		r.bufPos += n
		return n, nil
	}

	if r.body == nil {
		return 0, io.EOF
	}

	if len(p) > bufferSize {
		n, err = r.body.Read(p)
		r.offset += int64(n)
		return n, err
	}

	r.bufLen, err = r.body.Read(r.buffer)
	if err != nil && err != io.EOF {
		return 0, err
	}

	if r.bufLen > 0 {
		r.bufPos = 0
		n = copy(p, r.buffer[:r.bufLen])
		r.bufPos = n
		r.offset += int64(n)
		return n, nil
	}

	return 0, io.EOF
}

func (r *S3ReadSeeker) Seek(offset int64, whence int) (int64, error) {
	var newOffset int64
	switch whence {
	case io.SeekStart:
		newOffset = offset
	case io.SeekCurrent:
		newOffset = r.offset + offset
	case io.SeekEnd:
		newOffset = r.contentLen + offset
	default:
		return 0, fmt.Errorf("invalid whence")
	}

	if newOffset < 0 {
		return 0, fmt.Errorf("negative offset")
	}

	if newOffset >= r.contentLen {
		return r.contentLen, nil
	}

	if newOffset >= r.offset-int64(r.bufLen) && newOffset < r.offset {
		r.bufPos = int(newOffset - (r.offset - int64(r.bufLen)))
		r.offset = newOffset
		return newOffset, nil
	}

	if r.body != nil {
		r.body.Close()
		r.body = nil
	}
	r.bufLen = 0
	r.bufPos = 0

	result, err := r.client.GetObject(&s3.GetObjectInput{
		Bucket: aws.String(r.bucket),
		Key:    aws.String(r.key),
		Range:  aws.String(fmt.Sprintf("bytes=%d-", newOffset)),
	})
	if err != nil {
		return 0, err
	}

	r.body = result.Body
	r.offset = newOffset
	return newOffset, nil
}

func (r *S3ReadSeeker) Close() error {
	if r.body != nil {
		return r.body.Close()
	}
	return nil
}