package backend

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"

	"github.com/mem9-ai/drive9/pkg/datastore"
)

// stageMediaFile streams one stored object into an already-owned task
// workspace. It keeps large S3 objects out of heap memory and applies the
// caller's hard source cap while copying.
func (b *Dat9Backend) stageMediaFile(ctx context.Context, workspace *MediaExtractTaskWorkspace, f *datastore.File, name string, maxBytes int64) (string, int64, error) {
	if f.ContentLayout == datastore.ContentLayoutAppendLog {
		return "", 0, ErrAppendLogUnsupported
	}
	if maxBytes > 0 && f.SizeBytes > maxBytes {
		return "", 0, ErrMediaExtractSourceTooLarge
	}
	if f.StorageType == datastore.StorageDB9 {
		return workspace.StageReader(ctx, name, bytes.NewReader(f.ContentBlob), maxBytes)
	}
	if f.StorageType != datastore.StorageS3 {
		return "", 0, fmt.Errorf("unsupported storage type: %s", f.StorageType)
	}
	if b.s3 == nil {
		return "", 0, fmt.Errorf("s3 client not configured")
	}
	rc, err := b.s3.GetObject(ctx, f.StorageRef)
	if err != nil {
		return "", 0, err
	}
	defer func() { _ = rc.Close() }()
	return workspace.StageReader(ctx, name, rc, maxBytes)
}

func readBoundedFile(path string, maxBytes int64) ([]byte, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer func() { _ = f.Close() }()
	reader := io.Reader(f)
	if maxBytes > 0 {
		reader = io.LimitReader(f, maxBytes+1)
	}
	data, err := io.ReadAll(reader)
	if err != nil {
		return nil, err
	}
	if maxBytes > 0 && int64(len(data)) > maxBytes {
		return nil, ErrMediaExtractSourceTooLarge
	}
	return data, nil
}