package nextgentask

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
	"unsafe"

	sdkcommon "github.com/pingcap/metering_sdk/common"
	sdkconfig "github.com/pingcap/metering_sdk/config"
	"github.com/pingcap/metering_sdk/storage"
	metawriter "github.com/pingcap/metering_sdk/writer/meta"
	meteringwriter "github.com/pingcap/metering_sdk/writer/metering"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	deltalake "github.com/tidbcloud/metering/delta-rs-golang"
	"github.com/tidbcloud/metering/pkg/common/checkpoint"
	"github.com/tidbcloud/metering/pkg/common/config"
	"github.com/tidbcloud/metering/pkg/common/models"
	"github.com/tidbcloud/metering/pkg/common/utils"
	"github.com/tidbcloud/metering/pkg/nextgenmaintenance"
	"github.com/tidbcloud/metering/pkg/nextgenmetering"
	"github.com/tidbcloud/metering/pkg/nextgenrepoter"
	nextgentaskspkg "github.com/tidbcloud/metering/pkg/nextgentasks"
	"go.uber.org/zap"
)

const (
	logicalClusterID            = "logical-ft"
	sharedPoolID                = "sharedpool-ft"
	tenantIDValue               = uint64(60001)
	diskTypeValue               = "io2"
	diskIopsValue               = "4000"
	rowStorageBytes             = 1 * 1024 * 1024
	columnStorageBytes          = 5 * 1024 * 1024
	backupWrittenBytes          = 2 * 1024 * 1024
	olapRUValue                 = 321
	oltpRUValue                 = 654
	resourceManagerCrossAZBytes = 9 * 1024 * 1024
	proxyCrossZoneBytes         = 7 * 1024 * 1024
	proxyOutBoundBytes          = 3 * 1024 * 1024
)

var (
	initDeltaLockOnce sync.Once
	nextGenEnvCounter int64
)

func ensureDeltaLock() {
	initDeltaLockOnce.Do(func() {
		nextgentaskspkg.InitGlobalDeltaLakeLock(zap.NewNop())
	})
}

type memoryCheckpoint struct {
	mu   sync.Mutex
	data *checkpoint.CheckpointData
}

func newMemoryCheckpoint(initial *checkpoint.CheckpointData) *memoryCheckpoint {
	return &memoryCheckpoint{data: cloneCheckpointData(initial)}
}

func (m *memoryCheckpoint) GetCheckpoint(ctx context.Context, taskName string) (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.data == nil {
		return 0, nil
	}
	return m.data.Checkpoint, nil
}

func (m *memoryCheckpoint) SaveCheckpoint(ctx context.Context, taskName string, cp int64) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.data == nil {
		m.data = &checkpoint.CheckpointData{TaskName: taskName}
	}
	m.data.TaskName = taskName
	m.data.Checkpoint = cp
	m.data.UpdatedAt = time.Now().Unix()
	return nil
}

func (m *memoryCheckpoint) GetCheckpointData(ctx context.Context, taskName string) (*checkpoint.CheckpointData, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.data == nil {
		return &checkpoint.CheckpointData{TaskName: taskName}, nil
	}
	return cloneCheckpointData(m.data), nil
}

func (m *memoryCheckpoint) SaveCheckpointData(ctx context.Context, data *checkpoint.CheckpointData) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.data = cloneCheckpointData(data)
	return nil
}

func (m *memoryCheckpoint) Last() *checkpoint.CheckpointData {
	m.mu.Lock()
	defer m.mu.Unlock()
	return cloneCheckpointData(m.data)
}

func cloneCheckpointData(data *checkpoint.CheckpointData) *checkpoint.CheckpointData {
	if data == nil {
		return nil
	}
	cloned := *data
	return &cloned
}

type metaFixture struct {
	Metadata map[string]interface{} `json:"metadata"`
}

type meteringFixture struct {
	Timestamp    int64                    `json:"timestamp"`
	Category     string                   `json:"category"`
	SelfID       string                   `json:"self_id"`
	SharedPoolID string                   `json:"shared_pool_id"`
	Part         int                      `json:"part"`
	Data         []map[string]interface{} `json:"data"`
}

// testWriters holds metering and meta writers for test data
type testWriters struct {
	meteringWriter *meteringwriter.MeteringWriter
	metaWriter     *metawriter.MetaWriter
}

// newTestWriters creates writers for test data using the loader's storage provider
func newTestWriters(storageProvider storage.ObjectStorageProvider, log *zap.Logger) *testWriters {
	sdkCfg := sdkconfig.DefaultConfig()
	sdkCfg = sdkCfg.WithLogger(log)
	return &testWriters{
		meteringWriter: meteringwriter.NewMeteringWriter(storageProvider, sdkCfg),
		metaWriter:     metawriter.NewMetaWriter(storageProvider, sdkCfg),
	}
}

// WriteMeteringData writes metering data for testing
func (w *testWriters) WriteMeteringData(ctx context.Context, data *sdkcommon.MeteringData) error {
	return w.meteringWriter.Write(ctx, data)
}

// WriteMetaData writes metadata for testing
func (w *testWriters) WriteMetaData(ctx context.Context, clusterID string, metaType sdkcommon.MetaType, timestamp int64, data *sdkcommon.MetaData) error {
	// Set the cluster ID and meta type on the data
	data.ClusterID = clusterID
	data.Type = metaType
	data.ModifyTS = timestamp
	return w.metaWriter.Write(ctx, data)
}

// Close closes all writers
func (w *testWriters) Close() error {
	var errs []error
	if err := w.meteringWriter.Close(); err != nil {
		errs = append(errs, err)
	}
	if err := w.metaWriter.Close(); err != nil {
		errs = append(errs, err)
	}
	if len(errs) > 0 {
		return fmt.Errorf("errors closing writers: %v", errs)
	}
	return nil
}

type nextGenFTEnv struct {
	conf           *config.ServiceConfig
	startTimestamp int64
	records        []models.BillingReportV2Deltalake
	checkpoint     *memoryCheckpoint
}

func setupNextGenFTEnv(t *testing.T) nextGenFTEnv {
	counter := atomic.AddInt64(&nextGenEnvCounter, 1)
	base := time.Now().UTC().Truncate(time.Minute)
	start := base.Add(-10*time.Minute - time.Duration(counter)*time.Minute)
	return setupNextGenFTEnvAt(t, start)
}

func setupNextGenFTEnvAt(t *testing.T, startTime time.Time) nextGenFTEnv {
	t.Helper()
	require.NoError(t, os.RemoveAll("/tmp/nextgen-metering-data"))

	logger := zap.NewNop()
	ensureDeltaLock()

	startTime = startTime.UTC().Truncate(time.Minute)
	startTimestamp := startTime.Unix()
	delayMinutes := 9

	deltaDir := filepath.Join(t.TempDir(), "delta-table")
	// Create deltaDir to ensure it exists before delta table operations
	require.NoError(t, os.MkdirAll(deltaDir, 0755))
	cacheDir := filepath.Join(t.TempDir(), "cache")
	meterPrefix := filepath.Join("ft", fmt.Sprintf("nextgentask_%d", time.Now().UnixNano())) + "/"

	conf := &config.ServiceConfig{
		Vendor:   "local",
		Env:      "dev",
		LogLevel: "info",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			StartTime:      startTimestamp,
			DelayMinutes:   delayMinutes,
			DeltaTablePath: deltaDir,
			MeteringPrefix: meterPrefix,
			CachePath:      cacheDir,
			StorageOptions: map[string]string{},
		},
	}

	ctx := context.Background()

	loader := nextgenmetering.NewMeteringLoader(conf, logger)
	require.NoError(t, loader.Init())
	t.Cleanup(func() {
		_ = loader.Close()
	})
	t.Cleanup(func() {
		_ = os.RemoveAll(filepath.Join("/tmp/nextgen-metering-data", meterPrefix))
	})

	// Create writers for test data preparation
	writers := newTestWriters(loader.GetStorageProvider(), logger)
	t.Cleanup(func() {
		_ = writers.Close()
	})

	prepareMetadata(t, ctx, loader, writers, startTimestamp)
	prepareStorageSizeMetering(t, ctx, loader, writers, startTimestamp)
	prepareDFSStatsMetering(t, ctx, loader, writers, startTimestamp)
	prepareResourceManagerMetering(t, ctx, loader, writers, startTimestamp)
	prepareProxyMetering(t, ctx, loader, writers, startTimestamp)
	prepareDXFMetering(t, ctx, loader, writers, startTimestamp)
	prepareControlPlaneMetering(t, ctx, loader, writers, startTimestamp)

	cp := newMemoryCheckpoint(&checkpoint.CheckpointData{TaskName: "nextgen-metering"})

	task, err := nextgentaskspkg.NewNextGenTaskWithCheckpointManager(conf, logger, cp)
	require.NoError(t, err)

	err = task.Run(ctx)
	require.NoError(t, err)

	table, err := deltalake.OpenDeltaLakeTable(conf.NextGenMeteringConfig.DeltaTablePath, conf.NextGenMeteringConfig.StorageOptions)
	require.NoError(t, err)
	t.Cleanup(func() {
		table.Close()
	})

	rows, err := table.ReadData(100)
	require.NoError(t, err)
	require.NotEmpty(t, rows)

	records := make([]models.BillingReportV2Deltalake, 0, len(rows))
	for _, row := range rows {
		payload, err := json.Marshal(row)
		require.NoError(t, err)
		var rec models.BillingReportV2Deltalake
		require.NoError(t, json.Unmarshal(payload, &rec))
		records = append(records, rec)
	}

	return nextGenFTEnv{
		conf:           conf,
		startTimestamp: startTimestamp,
		records:        records,
		checkpoint:     cp,
	}
}

func TestNextGenTaskFunctional_MeteringSingleMinute(t *testing.T) {
	env := setupNextGenFTEnv(t)

	saved := env.checkpoint.Last()
	require.NotNil(t, saved)
	assert.Equal(t, env.startTimestamp+60, saved.Checkpoint)
	assert.Equal(t, env.conf.NextGenMeteringConfig.DeltaTablePath, saved.DeltaTablePath)
	assert.GreaterOrEqual(t, saved.DeltaTableVersion, int64(0))

	records := make(map[string]models.BillingReportV2Deltalake, len(env.records))
	for _, rec := range env.records {
		records[rec.Unit] = rec
	}

	taskRuUnit := fmt.Sprintf("%s:%s:%s:%s:%s", models.NextGenSKUTaskRU, env.conf.Vendor, "us-west-2", "premium", "import-into")

	expectedUnits := []string{
		fmt.Sprintf("%s:%s:%s:%s:%s:%s:%s:%s", models.NextGenSKUDiskStorageSize, env.conf.Vendor, "us-west-2", "premium", "row", "plus", "data", diskTypeValue),
		fmt.Sprintf("%s:%s:%s:%s:%s:%s:%s:%s", models.NextGenSKUDiskStorageSize, env.conf.Vendor, "us-west-2", "premium", "column", "plus", "data", diskTypeValue),
		fmt.Sprintf("%s:%s:%s:%s:%s", models.NextGenSKUBackupStorage, env.conf.Vendor, "us-west-2", "premium", "auto"),
		fmt.Sprintf("%s:%s:%s:%s:%s:total", models.NextGenSKURCU, env.conf.Vendor, "us-west-2", "premium", "all"), // Combined RCU (OLAP + OLTP)
		fmt.Sprintf("%s:%s:%s:%s:%s", models.NextGenSKUDataTransferCrossAZ, env.conf.Vendor, "us-west-2", "premium", models.NextGenSourceNamePD),
		fmt.Sprintf("%s:%s:%s:%s:%s", models.NextGenSKUDataTransferOut, env.conf.Vendor, "us-west-2", "premium", models.NextGenSourceNameProxy),
		fmt.Sprintf("%s:%s:%s:%s:%s", models.NextGenSKUDataTransferCrossAZ, env.conf.Vendor, "us-west-2", "premium", models.NextGenSourceNameProxy),
		taskRuUnit,
		// ControlPlane disk SKUs
		fmt.Sprintf("%s:%s:%s:%s:%s:%s:%s:%s", models.NextGenSKUDiskIOPS, env.conf.Vendor, "us-west-2", "premium", "row", "plus", "data", diskTypeValue),
		fmt.Sprintf("%s:%s:%s:%s:%s:%s:%s:%s", models.NextGenSKUDiskThroughput, env.conf.Vendor, "us-west-2", "premium", "row", "plus", "data", diskTypeValue),
		fmt.Sprintf("%s:%s:%s:%s:%s:%s:%s:%s", models.NextGenSKUDiskIOPS, env.conf.Vendor, "us-west-2", "premium", "row", "plus", "raft", diskTypeValue),
		fmt.Sprintf("%s:%s:%s:%s:%s:%s:%s:%s", models.NextGenSKUDiskThroughput, env.conf.Vendor, "us-west-2", "premium", "row", "plus", "raft", diskTypeValue),
	}

	require.Equal(t, len(expectedUnits), len(records))
	for _, unit := range expectedUnits {
		rec, ok := records[unit]
		require.Truef(t, ok, "unit %s not found in delta table", unit)
		assert.Equal(t, tenantIDValue, rec.TenantID)
		assert.Equal(t, env.startTimestamp, rec.CostTimestamp)

		// For ControlPlane disk SKUs, cost is the disk count (3), otherwise it's 1
		expectedCost := int64(1)
		if strings.HasPrefix(unit, models.NextGenSKUDiskIOPS+":") || strings.HasPrefix(unit, models.NextGenSKUDiskThroughput+":") {
			expectedCost = 3 // disk count from test data
		}
		assert.Equal(t, expectedCost, rec.Cost)

		assert.Equal(t, "v2", rec.SchemeVersion)
		assert.NotEmpty(t, rec.CostID)
		assert.True(t, rec.CreatedAt > 0)
		assert.Equal(t, utils.GenerateHourDimension(rec.CreatedAt), rec.HourDimension)

		labels := parseLabels(t, rec.Labels)
		expectedSKUType := string(models.NextGenSKUTypePublic)
		if unit == taskRuUnit {
			expectedSKUType = string(models.NextGenSKUTypePrivate)
		}
		// Throughput SKUs should be private
		if strings.HasPrefix(unit, models.NextGenSKUDiskThroughput+":") {
			expectedSKUType = string(models.NextGenSKUTypePrivate)
		}
		assert.Equal(t, expectedSKUType, rec.SKUType)

		switch unit {
		case fmt.Sprintf("%s:%s:%s:%s:%s:%s:%s:%s", models.NextGenSKUDiskStorageSize, env.conf.Vendor, "us-west-2", "premium", "row", "plus", "data", diskTypeValue):
			assert.Equal(t, fmt.Sprintf("%d", rowStorageBytes), rec.UnitValue)
			assert.Equal(t, logicalClusterID, labels[models.NextGenSKULabelKeyClusterID])
			assert.Equal(t, sharedPoolID, labels[models.NextGenSKULabelKeySharedPoolID])
			assert.Equal(t, "pd", labels[models.NextGenSKULabelKeySourceName])
			assert.Equal(t, "bytes", labels[models.NextGenSKULabelKeySourceUnit])
		case fmt.Sprintf("%s:%s:%s:%s:%s:%s:%s:%s", models.NextGenSKUDiskStorageSize, env.conf.Vendor, "us-west-2", "premium", "column", "plus", "data", diskTypeValue):
			assert.Equal(t, fmt.Sprintf("%d", columnStorageBytes), rec.UnitValue)
			assert.Equal(t, logicalClusterID, labels[models.NextGenSKULabelKeyClusterID])
			assert.Equal(t, sharedPoolID, labels[models.NextGenSKULabelKeySharedPoolID])
			assert.Equal(t, "pd", labels[models.NextGenSKULabelKeySourceName])
			assert.Equal(t, "bytes", labels[models.NextGenSKULabelKeySourceUnit])
		case fmt.Sprintf("%s:%s:%s:%s:%s", models.NextGenSKUBackupStorage, env.conf.Vendor, "us-west-2", "premium", "auto"):
			assert.Equal(t, fmt.Sprintf("%d", backupWrittenBytes), rec.UnitValue)
			assert.Equal(t, sharedPoolID, labels[models.NextGenSKULabelKeySharedPoolID])
			assert.Equal(t, "rfengine", labels[models.NextGenSKULabelKeySourceName])
			assert.Equal(t, "bytes", labels[models.NextGenSKULabelKeySourceUnit])
			_, hasCluster := labels[models.NextGenSKULabelKeyClusterID]
			assert.False(t, hasCluster)
		case fmt.Sprintf("%s:%s:%s:%s:%s:total", models.NextGenSKURCU, env.conf.Vendor, "us-west-2", "premium", "all"):
			// Combined RCU (OLAP + OLTP merged into "all")
			assert.Equal(t, "1", rec.UnitValue)
			totalRU := olapRUValue + oltpRUValue
			assert.Equal(t, fmt.Sprintf("%d", totalRU), labels[models.NextGenSKULabelKeyIntervalValue])
			assert.Equal(t, fmt.Sprintf("%d", olapRUValue), labels[models.NextGenSKULabelKeyOLAPRU])
			assert.Equal(t, fmt.Sprintf("%d", oltpRUValue), labels[models.NextGenSKULabelKeyOLTPRU])
			assert.Equal(t, logicalClusterID, labels[models.NextGenSKULabelKeyClusterID])
			assert.Equal(t, sharedPoolID, labels[models.NextGenSKULabelKeySharedPoolID])
			assert.Equal(t, "ru", labels[models.NextGenSKULabelKeySourceUnit])
			assert.Equal(t, models.NextgenFixMinRU, labels[models.NextGenSKULabelKeyIntervalMin])
		case fmt.Sprintf("%s:%s:%s:%s:%s", models.NextGenSKUDataTransferCrossAZ, env.conf.Vendor, "us-west-2", "premium", models.NextGenSourceNamePD):
			assert.Equal(t, fmt.Sprintf("%d", resourceManagerCrossAZBytes), rec.UnitValue)
			assert.Equal(t, logicalClusterID, labels[models.NextGenSKULabelKeyClusterID])
			assert.Equal(t, sharedPoolID, labels[models.NextGenSKULabelKeySharedPoolID])
			assert.Equal(t, "bytes", labels[models.NextGenSKULabelKeySourceUnit])
		case fmt.Sprintf("%s:%s:%s:%s:%s", models.NextGenSKUDataTransferOut, env.conf.Vendor, "us-west-2", "premium", models.NextGenSourceNameProxy):
			assert.Equal(t, fmt.Sprintf("%d", proxyOutBoundBytes), rec.UnitValue)
			assert.Equal(t, logicalClusterID, labels[models.NextGenSKULabelKeyClusterID])
			assert.Equal(t, sharedPoolID, labels[models.NextGenSKULabelKeySharedPoolID])
			assert.Equal(t, "bytes", labels[models.NextGenSKULabelKeySourceUnit])
		case fmt.Sprintf("%s:%s:%s:%s:%s", models.NextGenSKUDataTransferCrossAZ, env.conf.Vendor, "us-west-2", "premium", models.NextGenSourceNameProxy):
			assert.Equal(t, fmt.Sprintf("%d", proxyCrossZoneBytes), rec.UnitValue)
			assert.Equal(t, logicalClusterID, labels[models.NextGenSKULabelKeyClusterID])
			assert.Equal(t, sharedPoolID, labels[models.NextGenSKULabelKeySharedPoolID])
			assert.Equal(t, "bytes", labels[models.NextGenSKULabelKeySourceUnit])
		case taskRuUnit:
			assert.Equal(t, "1", rec.UnitValue)
			assert.Equal(t, int64(1), rec.Cost)
			assert.Equal(t, logicalClusterID, labels[models.NextGenSKULabelKeyClusterID])
			assert.Equal(t, sharedPoolID, labels[models.NextGenSKULabelKeySharedPoolID])
			assert.Equal(t, "30001", labels[models.NextGenSKULabelKeyTaskID])
			assert.Equal(t, "import-into", labels[models.NextGenSKULabelKeyTaskType])
			assert.Equal(t, "1", labels[models.NextGenSKULabelKeyVersion])
			assert.Equal(t, "8", labels[models.NextGenSKULabelKeyGetExtReq])
			assert.Equal(t, "24", labels[models.NextGenSKULabelKeyPutExtReq])
			assert.Equal(t, "109689961", labels[models.NextGenSKULabelKeyReadBytes])
			assert.Equal(t, "37004046", labels[models.NextGenSKULabelKeyWriteBytes])
		}
	}
}

func TestNextGenReporterFunctional_ReportConversion(t *testing.T) {
	env := setupNextGenFTEnv(t)

	reports := make([]*models.BillingReportV2Deltalake, 0, len(env.records))
	var maxCreatedAt int64
	for i := range env.records {
		recCopy := env.records[i]
		if recCopy.CreatedAt > maxCreatedAt {
			maxCreatedAt = recCopy.CreatedAt
		}
		reports = append(reports, &recCopy)
	}

	loader := &ftReporterLoader{
		checkpoint: env.startTimestamp - 60,
		reports:    reports,
	}

	sender := &spyMeasureSender{}

	reporter := &nextgenrepoter.NextGenBillingRepoter{}
	setUnexportedField(t, reporter, "conf", env.conf)
	setUnexportedField(t, reporter, "loader", loader)
	setUnexportedField(t, reporter, "log", zap.NewNop())
	reporter.SetSender(sender)

	ctx := context.Background()
	require.NoError(t, reporter.Report(ctx))

	require.Equal(t, len(env.records), len(sender.measures))
	assert.Equal(t, env.startTimestamp-60, loader.loadStart)
	assert.GreaterOrEqual(t, loader.loadEnd, env.startTimestamp)
	assert.Equal(t, maxCreatedAt, loader.savedCheckpoint)

	measuresByUnit := make(map[string]models.BillingReport, len(sender.measures))
	for _, measure := range sender.measures {
		measuresByUnit[measure.Unit] = measure
		assert.Equal(t, env.startTimestamp, measure.CostTimestamp)
	}

	const bytesPerMiB int64 = 1024 * 1024
	type bytesExpectation struct {
		bytes int64
	}

	expectedByteUnits := map[string]bytesExpectation{
		fmt.Sprintf("%s:%s:%s:%s:%s:%s:%s:%s", models.NextGenSKUDiskStorageSize, env.conf.Vendor, "us-west-2", "premium", "row", "plus", "data", diskTypeValue):    {bytes: int64(rowStorageBytes)},
		fmt.Sprintf("%s:%s:%s:%s:%s:%s:%s:%s", models.NextGenSKUDiskStorageSize, env.conf.Vendor, "us-west-2", "premium", "column", "plus", "data", diskTypeValue): {bytes: int64(columnStorageBytes)},
		fmt.Sprintf("%s:%s:%s:%s:%s", models.NextGenSKUBackupStorage, env.conf.Vendor, "us-west-2", "premium", "auto"):                                             {bytes: int64(backupWrittenBytes)},
		fmt.Sprintf("%s:%s:%s:%s:%s", models.NextGenSKUDataTransferCrossAZ, env.conf.Vendor, "us-west-2", "premium", models.NextGenSourceNamePD):                   {bytes: int64(resourceManagerCrossAZBytes)},
		fmt.Sprintf("%s:%s:%s:%s:%s", models.NextGenSKUDataTransferOut, env.conf.Vendor, "us-west-2", "premium", models.NextGenSourceNameProxy):                    {bytes: int64(proxyOutBoundBytes)},
		fmt.Sprintf("%s:%s:%s:%s:%s", models.NextGenSKUDataTransferCrossAZ, env.conf.Vendor, "us-west-2", "premium", models.NextGenSourceNameProxy):                {bytes: int64(proxyCrossZoneBytes)},
	}

	for unit, expectation := range expectedByteUnits {
		measure, ok := measuresByUnit[unit]
		require.Truef(t, ok, "expected measure for unit %s", unit)
		expectedMiB := (expectation.bytes + bytesPerMiB - 1) / bytesPerMiB
		assert.Equalf(t, expectedMiB, measure.Cost, "unit %s cost mismatch", unit)
		labels := parseLabels(t, measure.Labels)
		assert.Equalf(t, fmt.Sprintf("%d", expectation.bytes), labels[models.NextGenSKULabelKeyOriginalBytes], "unit %s original bytes mismatch", unit)
	}

	unitCostOne := []string{
		fmt.Sprintf("%s:%s:%s:%s:%s:total", models.NextGenSKURCU, env.conf.Vendor, "us-west-2", "premium", "all"), // Combined RCU
	}

	for _, unit := range unitCostOne {
		measure, ok := measuresByUnit[unit]
		require.Truef(t, ok, "expected measure for unit %s", unit)
		assert.Equalf(t, int64(1), measure.Cost, "unit %s cost mismatch", unit)
		labels := parseLabels(t, measure.Labels)
		_, hasOriginalBytes := labels[models.NextGenSKULabelKeyOriginalBytes]
		assert.Falsef(t, hasOriginalBytes, "unit %s should not contain original bytes label", unit)
	}
}

func TestNextGenMaintenanceTaskFunctional_DeleteOldPartitions(t *testing.T) {
	env := setupNextGenFTEnv(t)
	require.NotEmpty(t, env.records)

	preTable, err := deltalake.OpenDeltaLakeTable(env.conf.NextGenMeteringConfig.DeltaTablePath, env.conf.NextGenMeteringConfig.StorageOptions)
	require.NoError(t, err)
	rowsBefore, err := preTable.ReadData(100)
	require.NoError(t, err)
	preTable.Close()
	require.NotEmpty(t, rowsBefore)

	cutoff := env.records[0].CreatedAt + 3600
	futureNow := time.Unix(cutoff+1800, 0).UTC()

	env.conf.NextGenMeteringConfig.MaintenanceConfig = &config.NextGenMaintenanceConfig{
		Enable:                true,
		EnableDelete:          true,
		DeleteBeforeTimestamp: cutoff,
		EnableVacuum:          false,
	}

	task, err := nextgentaskspkg.NewNextGenMaintenanceTask(env.conf, zap.NewNop())
	require.NoError(t, err)

	customRunner, err := nextgenmaintenance.NewNextGenMaintenance(env.conf, zap.NewNop())
	require.NoError(t, err)
	setUnexportedField(t, customRunner, "now", func() time.Time { return futureNow })
	setUnexportedField(t, task, "runner", customRunner)

	ctx := context.Background()
	require.NoError(t, task.Run(ctx))

	refreshed, err := deltalake.OpenDeltaLakeTable(env.conf.NextGenMeteringConfig.DeltaTablePath, env.conf.NextGenMeteringConfig.StorageOptions)
	require.NoError(t, err)
	t.Cleanup(func() { refreshed.Close() })

	rowsAfter, err := refreshed.ReadData(100)
	require.NoError(t, err)
	assert.Empty(t, rowsAfter)
}

func prepareResourceManagerMetering(t *testing.T, ctx context.Context, loader *nextgenmetering.MeteringLoader, writers *testWriters, timestamp int64) {
	t.Helper()
	path := fixturePath("resource-manager", "pd13091331099564961501-0.json")
	data := loadMeteringFixture(t, path)

	record := deepCopyMap(data.Data[0])
	record[models.NextGenMeteringLogicalClusterID] = logicalClusterID
	record[models.NextGenMeteringSourceName] = models.NextGenSourceNamePD
	record[models.NextGenMeteringOLAPRU] = map[string]interface{}{"value": float64(olapRUValue), "unit": "RU"}
	record[models.NextGenMeteringOLTPRU] = map[string]interface{}{"value": float64(oltpRUValue), "unit": "RU"}
	setMeteringValue(record, models.NextGenMeteringCrossAZTraffic, resourceManagerCrossAZBytes)

	writeMetering(t, ctx, loader, writers, meteringFixture{
		Timestamp:    timestamp,
		Category:     models.NextGenCategoryResourceManager,
		SelfID:       data.SelfID,
		SharedPoolID: sharedPoolID,
		Part:         data.Part,
		Data:         []map[string]interface{}{record},
	})
}

func prepareProxyMetering(t *testing.T, ctx context.Context, loader *nextgenmetering.MeteringLoader, writers *testWriters, timestamp int64) {
	t.Helper()
	path := fixturePath("proxy", "8232a821_1721_4221_a2ae_ece91639151d-0.json")
	data := loadMeteringFixture(t, path)

	record := deepCopyMap(data.Data[0])
	record[models.NextGenMeteringLogicalClusterID] = logicalClusterID
	record[models.NextGenMeteringSourceName] = models.NextGenSourceNameProxy
	setMeteringValue(record, models.NextGenMeteringCrossZoneBytes, proxyCrossZoneBytes)
	setMeteringValue(record, models.NextGenMeteringOutBoundBytes, proxyOutBoundBytes)

	writeMetering(t, ctx, loader, writers, meteringFixture{
		Timestamp:    timestamp,
		Category:     models.NextGenCategoryProxy,
		SelfID:       data.SelfID,
		SharedPoolID: sharedPoolID,
		Part:         data.Part,
		Data:         []map[string]interface{}{record},
	})
}

func prepareDXFMetering(t *testing.T, ctx context.Context, loader *nextgenmetering.MeteringLoader, writers *testWriters, timestamp int64) {
	t.Helper()
	path := fixturePath("dxf", "b5367e66_f0be_4c8d_89f2_653950767d8e-0.json")
	data := loadMeteringFixture(t, path)

	record := deepCopyMap(data.Data[0])
	record[models.NextGenMeteringLogicalClusterID] = logicalClusterID
	record[models.NextGenMeteringSourceName] = models.NextGenSourceNameDXF
	record[models.NextGenMeteringTaskType] = "import-into"
	record[models.NextGenMeteringTaskID] = 30001
	record[models.NextGenMeteringGetExtRequests] = 8
	record[models.NextGenMeteringPutExtRequests] = 24
	setMeteringValue(record, models.NextGenMeteringReadDataBytes, int(109689961))
	setMeteringValue(record, models.NextGenMeteringWriteDataBytes, int(37004046))

	writeMetering(t, ctx, loader, writers, meteringFixture{
		Timestamp:    timestamp,
		Category:     models.NextGenCategoryDXF,
		SelfID:       data.SelfID,
		SharedPoolID: sharedPoolID,
		Part:         data.Part,
		Data:         []map[string]interface{}{record},
	})
}

func prepareMetadata(t *testing.T, ctx context.Context, loader *nextgenmetering.MeteringLoader, writers *testWriters, timestamp int64) {
	t.Helper()

	metaPath := fixturePath("meta", "1759046592.json")
	raw, err := os.ReadFile(metaPath)
	require.NoError(t, err)

	var fixture metaFixture
	require.NoError(t, json.Unmarshal(raw, &fixture))

	logicMetadata := deepCopyMap(fixture.Metadata)
	logicMetadata[models.NextGenMetaTenantID] = fmt.Sprintf("%d", tenantIDValue)
	logicMetadata[models.NextGenMeteringRegion] = "us-west-2"
	logicMetadata[models.NextGenMetaDiskType] = diskTypeValue
	logicMetadata[models.NextGenMetaDiskIOPS] = diskIopsValue
	logicMetadata["physicalID"] = logicalClusterID
	logicMetadata["plan"] = "premium"

	logic := loader.CreateMetaData(logicalClusterID, timestamp, logicMetadata)
	require.NoError(t, writers.WriteMetaData(ctx, logicalClusterID, sdkcommon.MetaTypeLogic, timestamp, logic))

	sharedMetadata := deepCopyMap(fixture.Metadata)
	sharedMetadata[models.NextGenMetaTenantID] = fmt.Sprintf("%d", tenantIDValue)
	sharedMetadata[models.NextGenMeteringRegion] = "us-west-2"
	sharedMetadata["physicalID"] = sharedPoolID
	shared := loader.CreateMetaData(sharedPoolID, timestamp, sharedMetadata)
	require.NoError(t, writers.WriteMetaData(ctx, sharedPoolID, sdkcommon.MetaTypeSharedpool, timestamp, shared))
}

func prepareStorageSizeMetering(t *testing.T, ctx context.Context, loader *nextgenmetering.MeteringLoader, writers *testWriters, timestamp int64) {
	t.Helper()
	path := fixturePath("storage-size", "pd13091331099564961501-0.json")
	data := loadMeteringFixture(t, path)

	record := deepCopyMap(data.Data[0])
	record[models.NextGenMeteringLogicalClusterID] = logicalClusterID
	record[models.NextGenMeteringSourceName] = "pd"
	record[models.NextGenMeteringVersion] = "1"
	setMeteringValue(record, models.NextGenMeteringRowStorageSize, rowStorageBytes)
	setMeteringValue(record, models.NextGenMeteringColumnStorageSize, columnStorageBytes)

	writeMetering(t, ctx, loader, writers, meteringFixture{
		Timestamp:    timestamp,
		Category:     models.NextGenCategoryStorageSize,
		SelfID:       data.SelfID,
		SharedPoolID: sharedPoolID,
		Part:         data.Part,
		Data:         []map[string]interface{}{record},
	})
}

func prepareDFSStatsMetering(t *testing.T, ctx context.Context, loader *nextgenmetering.MeteringLoader, writers *testWriters, timestamp int64) {
	t.Helper()
	path := fixturePath("dfs-stats", "pd13091331099564961501-0.json")
	data := loadMeteringFixture(t, path)

	record := deepCopyMap(data.Data[0])
	record[models.NextGenMeteringSourceName] = "rfengine"
	setMeteringValue(record, models.NextGenMeteringWrittenBytes, backupWrittenBytes)
	setMeteringValue(record, models.NextGenMeteringWriteRequests, 1)

	writeMetering(t, ctx, loader, writers, meteringFixture{
		Timestamp:    timestamp,
		Category:     models.NextGenCategoryDFSStats,
		SelfID:       data.SelfID,
		SharedPoolID: sharedPoolID,
		Part:         data.Part,
		Data:         []map[string]interface{}{record},
	})
}

func prepareControlPlaneMetering(t *testing.T, ctx context.Context, loader *nextgenmetering.MeteringLoader, writers *testWriters, timestamp int64) {
	t.Helper()
	path := fixturePath("control-plane", "control-plane-ft-0.json")
	data := loadMeteringFixture(t, path)

	record := deepCopyMap(data.Data[0])
	record[models.NextGenMeteringLogicalClusterID] = logicalClusterID
	record[models.NextGenMeteringSourceName] = "control-plane"

	writeMetering(t, ctx, loader, writers, meteringFixture{
		Timestamp:    timestamp,
		Category:     models.NextGenCategoryControlPlane,
		SelfID:       data.SelfID,
		SharedPoolID: sharedPoolID,
		Part:         data.Part,
		Data:         []map[string]interface{}{record},
	})
}

func setMeteringValue(target map[string]interface{}, key string, value int) {
	if mv, ok := target[key].(map[string]interface{}); ok {
		mv["value"] = float64(value)
		if _, unitExists := mv["unit"]; !unitExists {
			mv["unit"] = "Bytes"
		}
	} else {
		target[key] = map[string]interface{}{"value": float64(value), "unit": "Bytes"}
	}
}

func writeMetering(t *testing.T, ctx context.Context, loader *nextgenmetering.MeteringLoader, writers *testWriters, fixture meteringFixture) {
	meter := loader.CreateMeteringData(fixture.SelfID, fixture.Category, fixture.Timestamp, fixture.Data)
	meter.SharedPoolID = fixture.SharedPoolID
	require.NoError(t, writers.WriteMeteringData(ctx, meter))
}

func loadMeteringFixture(t *testing.T, path string) meteringFixture {
	t.Helper()
	raw, err := os.ReadFile(path)
	require.NoError(t, err)
	var fixture meteringFixture
	require.NoError(t, json.Unmarshal(raw, &fixture))
	return fixture
}

func deepCopyMap(input map[string]interface{}) map[string]interface{} {
	copied := make(map[string]interface{}, len(input))
	for k, v := range input {
		copied[k] = deepCopyValue(v)
	}
	return copied
}

func deepCopyValue(v interface{}) interface{} {
	switch val := v.(type) {
	case map[string]interface{}:
		return deepCopyMap(val)
	case []interface{}:
		result := make([]interface{}, len(val))
		for i, item := range val {
			result[i] = deepCopyValue(item)
		}
		return result
	default:
		return val
	}
}

func parseLabels(t *testing.T, labels string) map[string]string {
	t.Helper()
	var raw map[string]interface{}
	require.NoError(t, json.Unmarshal([]byte(labels), &raw))
	normalized := make(map[string]string, len(raw))
	for k, v := range raw {
		normalized[k] = fmt.Sprint(v)
	}
	return normalized
}

func fixturePath(parts ...string) string {
	base := []string{"..", "..", "..", "testdata", "nextgenfille"}
	base = append(base, parts...)
	return filepath.Join(base...)
}

func setUnexportedField(t *testing.T, target interface{}, fieldName string, value interface{}) {
	t.Helper()
	v := reflect.ValueOf(target)
	require.Equalf(t, reflect.Pointer, v.Kind(), "target must be a pointer, got %s", v.Kind())
	elem := v.Elem()
	field := elem.FieldByName(fieldName)
	require.Truef(t, field.IsValid(), "field %s not found", fieldName)

	valueV := reflect.ValueOf(value)
	// Special handling for interface types - check if value implements the interface
	if field.Type().Kind() == reflect.Interface {
		require.Truef(t, valueV.Type().Implements(field.Type()), "value type %s does not implement interface %s", valueV.Type(), field.Type())
	} else {
		require.Truef(t, valueV.Type().AssignableTo(field.Type()), "value type %s cannot be assigned to field %s of type %s", valueV.Type(), fieldName, field.Type())
	}

	reflect.NewAt(field.Type(), unsafe.Pointer(field.UnsafeAddr())).Elem().Set(valueV)
}

type ftReporterLoader struct {
	checkpoint      int64
	reports         []*models.BillingReportV2Deltalake
	savedCheckpoint int64
	loadStart       int64
	loadEnd         int64
}

func (l *ftReporterLoader) LoadMeteringDataInTimeRange(_ context.Context, startTime, endTime int64) ([]*models.BillingReportV2Deltalake, error) {
	if l.loadStart == 0 || startTime < l.loadStart {
		l.loadStart = startTime
	}
	if endTime > l.loadEnd {
		l.loadEnd = endTime
	}

	result := make([]*models.BillingReportV2Deltalake, 0, len(l.reports))
	for _, report := range l.reports {
		if report.CreatedAt >= startTime && (endTime == 0 || report.CreatedAt <= endTime) {
			result = append(result, report)
		}
	}
	return result, nil
}

func (l *ftReporterLoader) LoadCheckpoint(context.Context) (int64, error) {
	return l.checkpoint, nil
}

func (l *ftReporterLoader) SaveCheckpoint(_ context.Context, checkpoint int64) error {
	l.savedCheckpoint = checkpoint
	return nil
}

func (l *ftReporterLoader) LoadMeteringDataInTimeRangeWithDeltaRS(ctx context.Context, startTime, endTime int64) ([]*models.BillingReportV2Deltalake, error) {
	// Use the same implementation as LoadMeteringDataInTimeRange for testing purposes
	return l.LoadMeteringDataInTimeRange(ctx, startTime, endTime)
}

type spyMeasureSender struct {
	measures   []models.BillingReport
	initCalled bool
}

func (s *spyMeasureSender) InitMeasureHandler(string, string, bool, bool, ...string) error {
	s.initCalled = true
	return nil
}

func (s *spyMeasureSender) SendMeasure(_ context.Context, billingMeasure models.BillingReport) error {
	s.measures = append(s.measures, billingMeasure)
	return nil
}

func (s *spyMeasureSender) CheckMeasure(models.BillingReport) error {
	return nil
}