package vmquery

import (
	"context"
	"io"
	"net/http"
	"time"

	"github.com/tidbcloud/observability/vm-migration/internal/schema"
)

// ExportOptions controls one ExportShard invocation.
type ExportOptions struct {
	SourceURLTemplate string        // vmselect URL with {tenant} placeholder
	Match             string        // PromQL series selector; empty → MatchAll
	Timeout           time.Duration // per-request timeout
	RateLimit         int           // bytes/sec; 0 = unlimited
}

// ShardResult captures the outcome of ExportShard. Error empty = success.
type ShardResult struct {
	ShardID string
	Bytes   int64
	DurMs   int64
	Error   string
}

// ExportShard POSTs to /api/v1/export/native for one shard and counts response
// bytes. DurMs is measured from request start to body EOF. On ctx cancellation,
// returns the result with partial byte count + Error set.
func ExportShard(ctx context.Context, client *http.Client, shard schema.Shard, opts ExportOptions) ShardResult {
	started := time.Now()
	res := ShardResult{ShardID: shard.ShardID}
	if opts.Match == "" {
		opts.Match = MatchAll
	}
	byteCount, err := exportShardBytes(ctx, client, shard, opts.SourceURLTemplate, opts.Match, opts.Timeout, opts.RateLimit)
	res.Bytes = byteCount
	res.DurMs = time.Since(started).Milliseconds()
	if err != nil {
		res.Error = err.Error()
	}
	return res
}

// exportShardBytes is the byte-counting core. Kept private; callers use ExportShard.
func exportShardBytes(ctx context.Context, client *http.Client, shard schema.Shard, sourceTemplate, match string, timeout time.Duration, rateLimit int) (int64, error) {
	tenantURL := RenderURL(sourceTemplate, shard.Tenant)
	target := Endpoint(tenantURL, "/api/v1/export/native")
	req, cancel, err := NewFormRequest(ctx, target, map[string]string{
		"match[]": match,
		"start":   shard.TimeStart,
		"end":     shard.TimeEnd,
	}, timeout)
	if err != nil {
		return 0, err
	}
	defer cancel()
	resp, err := client.Do(req)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	if err := CheckResponse(resp); err != nil {
		return 0, err
	}

	var total int64
	buf := make([]byte, 65536)
	streamStart := time.Now()
	for {
		n, readErr := resp.Body.Read(buf)
		if n > 0 {
			total += int64(n)
			if rateLimit > 0 {
				expected := time.Duration(float64(total) / float64(rateLimit) * float64(time.Second))
				if elapsed := time.Since(streamStart); elapsed < expected {
					select {
					case <-time.After(expected - elapsed):
					case <-req.Context().Done():
						return total, req.Context().Err()
					}
				}
			}
		}
		if readErr == io.EOF {
			break
		}
		if readErr != nil {
			return total, readErr
		}
	}
	return total, nil
}