package statistic

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/tidbcloud/observability/o11ymetering/pkg/facade/iface"
	"go.uber.org/mock/gomock"

	mock_facade "github.com/tidbcloud/observability/o11ymetering/mock/pkg/facade"
	"github.com/tidbcloud/observability/o11ymetering/pkg/common/config"
	"github.com/tidbcloud/observability/o11ymetering/pkg/common/constants"
	"github.com/tidbcloud/observability/o11ymetering/pkg/common/message"
	"github.com/tidbcloud/observability/o11ymetering/pkg/common/models"
	"github.com/tidbcloud/observability/o11ymetering/pkg/facade"
	"github.com/tidbcloud/observability/o11ymetering/pkg/service/statistic/handler"
	"github.com/tidbcloud/observability/pkg/common/configs"
	"github.com/tidbcloud/observability/pkg/common/structs"
)

// setupStatisticMocks initializes gomock controller and mocks for statistic related interfaces.
func setupStatisticMocks(t *testing.T) (*gomock.Controller, *mock_facade.MockStatisticReaderWriter, *mock_facade.MockStatisticMergeReaderWriter, *mock_facade.MockStatisticReportReaderWriter, *mock_facade.MockScheduleTaskDataSync, *mock_facade.MockClient) {
	ctrl := gomock.NewController(t)
	mockStatisticRW := mock_facade.NewMockStatisticReaderWriter(ctrl)
	mockStatisticMergeRW := mock_facade.NewMockStatisticMergeReaderWriter(ctrl)
	mockStatisticReportRW := mock_facade.NewMockStatisticReportReaderWriter(ctrl)
	mockScheduleTaskDataSync := mock_facade.NewMockScheduleTaskDataSync(ctrl) // Assuming this mock exists
	mockS3Client := mock_facade.NewMockClient(ctrl)                           // Assuming this mock exists

	facade.SetStatisticReaderWriter(mockStatisticRW)
	facade.SetStatisticMergeReaderWriter(mockStatisticMergeRW)
	facade.SetStatisticReportReaderWriter(mockStatisticReportRW)
	facade.SetScheduleTaskDataSync(mockScheduleTaskDataSync)
	facade.SetS3Client(mockS3Client)

	return ctrl, mockStatisticRW, mockStatisticMergeRW, mockStatisticReportRW, mockScheduleTaskDataSync, mockS3Client
}

func TestReport(t *testing.T) {
	configs.Init(structs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})
	t.Run("Report success", func(t *testing.T) {
		ctrl, _, mockStatisticMergeRW, mockStatisticReportRW, mockScheduleTaskDataSync, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		req := message.GenerateStatisticReportReq{Timestamp: int32(time.Now().Unix())}
		period := time.Unix(int64(req.Timestamp), 0).Format("2006-01-02 15:04:05")

		mockScheduleTaskDataSync.EXPECT().LocalDataFilePath(gomock.Any(), constants.StatisticMergeTaskName, gomock.Any()).Return("mock/path/file")
		mockScheduleTaskDataSync.EXPECT().Download(gomock.Any(), constants.StatisticMergeTaskName, gomock.Any()).Return("mock/path/file", nil)

		// Mock SwitchStatisticMergeReaderWriter
		mockStatisticMergeRW.EXPECT().LatestMergedRecordPeriod(gomock.Any(), gomock.Any()).Return(period, nil).AnyTimes()
		mockStatisticMergeRW.EXPECT().MonthlyStatisticSum(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(100.0, nil).AnyTimes()
		mockStatisticMergeRW.EXPECT().MonthlyTenantStatistic(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]*models.MonthlyTenantStatistic{
			{Tenant: "t1", Env: constants.Prod, MeteringGroup: constants.DedicatedDataplane, SubCategory: "sub1", Category: constants.Metrics, SumValue: 50.0},
		}, nil).AnyTimes()

		mockStatisticReportRW.EXPECT().GetMonthlyCost(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(200.0, nil).AnyTimes()
		mockStatisticReportRW.EXPECT().InsertReport(gomock.Any(), gomock.Any()).Return(nil)

		originalGetCategoryShareProportion := config.GetCategoryShareProportion
		getCategoryShareProportion = func(category constants.Category, cloudService constants.CloudService) float64 {
			return 0.5
		}
		defer func() { getCategoryShareProportion = originalGetCategoryShareProportion }()

		original := switchStatisticMergeReaderWriter
		switchStatisticMergeReaderWriter = func(ctx context.Context, timestamp int64) iface.StatisticMergeReaderWriter {
			return mockStatisticMergeRW
		}
		defer func() {
			switchStatisticMergeReaderWriter = original
		}()
		resp, err := Report(context.Background(), req)
		assert.NoError(t, err)
		assert.Equal(t, message.GenerateStatisticReportResp{}, resp)
	})

	t.Run("Report ShareCostsStatisticReport error", func(t *testing.T) {
		ctrl, _, mockStatisticMergeRW, mockStatisticReportRW, mockScheduleTaskDataSync, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		req := message.GenerateStatisticReportReq{Timestamp: int32(time.Now().Unix())}
		period := time.Unix(int64(req.Timestamp), 0).Format("2006-01-02 15:04:05")

		mockScheduleTaskDataSync.EXPECT().LocalDataFilePath(gomock.Any(), constants.StatisticMergeTaskName, gomock.Any()).Return("mock/path/file")
		mockScheduleTaskDataSync.EXPECT().Download(gomock.Any(), constants.StatisticMergeTaskName, gomock.Any()).Return("mock/path/file", nil)

		mockStatisticMergeRW.EXPECT().LatestMergedRecordPeriod(gomock.Any(), gomock.Any()).Return(period, nil).AnyTimes()                                                // Mocked for SwitchStatisticMergeReaderWriter
		mockStatisticReportRW.EXPECT().GetMonthlyCost(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(0.0, errors.New("get monthly cost error")).Times(1) // Only first share request fails

		original := switchStatisticMergeReaderWriter
		switchStatisticMergeReaderWriter = func(ctx context.Context, timestamp int64) iface.StatisticMergeReaderWriter {
			return mockStatisticMergeRW
		}
		defer func() {
			switchStatisticMergeReaderWriter = original
		}()
		_, err := Report(context.Background(), req)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "get monthly cost error")
	})

	t.Run("Report InsertReport error", func(t *testing.T) {
		ctrl, _, mockStatisticMergeRW, mockStatisticReportRW, mockScheduleTaskDataSync, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		req := message.GenerateStatisticReportReq{Timestamp: int32(time.Now().Unix())}
		period := time.Unix(int64(req.Timestamp), 0).Format("2006-01-02 15:04:05")

		mockScheduleTaskDataSync.EXPECT().LocalDataFilePath(gomock.Any(), constants.StatisticMergeTaskName, gomock.Any()).Return("mock/path/file")
		mockScheduleTaskDataSync.EXPECT().Download(gomock.Any(), constants.StatisticMergeTaskName, gomock.Any()).Return("mock/path/file", nil)

		mockStatisticMergeRW.EXPECT().LatestMergedRecordPeriod(gomock.Any(), gomock.Any()).Return(period, nil).AnyTimes()
		mockStatisticMergeRW.EXPECT().MonthlyStatisticSum(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(100.0, nil).AnyTimes()
		mockStatisticMergeRW.EXPECT().MonthlyTenantStatistic(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]*models.MonthlyTenantStatistic{
			{Tenant: "t1", Env: constants.Prod, MeteringGroup: constants.DedicatedDataplane, SubCategory: "sub1", Category: constants.Metrics, SumValue: 50.0},
		}, nil).AnyTimes()

		mockStatisticReportRW.EXPECT().GetMonthlyCost(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(200.0, nil).AnyTimes()
		mockStatisticReportRW.EXPECT().InsertReport(gomock.Any(), gomock.Any()).Return(errors.New("insert report error"))

		originalGetCategoryShareProportion := config.GetCategoryShareProportion
		getCategoryShareProportion = func(category constants.Category, cloudService constants.CloudService) float64 {
			return 0.5
		}
		defer func() { getCategoryShareProportion = originalGetCategoryShareProportion }()

		original := switchStatisticMergeReaderWriter
		switchStatisticMergeReaderWriter = func(ctx context.Context, timestamp int64) iface.StatisticMergeReaderWriter {
			return mockStatisticMergeRW
		}
		defer func() {
			switchStatisticMergeReaderWriter = original
		}()
		_, err := Report(context.Background(), req)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "insert report error")
	})
}

func TestShareCostsStatisticReport(t *testing.T) {
	configs.Init(structs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})
	req := message.ShareCostsStatisticReportReq{
		Timestamp:     int32(time.Now().Unix()),
		StartPeriod:   "2023-01-01 00:00:00",
		EndPeriod:     "2023-02-01 00:00:00",
		StatisticName: constants.MetricsTimeseriesCount,
		Category:      constants.Metrics,
		CloudService:  constants.ElasticCompute,
	}

	t.Run("ShareCostsStatisticReport success", func(t *testing.T) {
		ctrl, _, mockStatisticMergeRW, mockStatisticReportRW, _, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		mockStatisticReportRW.EXPECT().GetMonthlyCost(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(100.0, nil)
		mockStatisticMergeRW.EXPECT().MonthlyStatisticSum(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(50.0, nil)
		mockStatisticMergeRW.EXPECT().MonthlyTenantStatistic(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]*models.MonthlyTenantStatistic{
			{Tenant: "t1", Env: constants.Prod, MeteringGroup: constants.DedicatedDataplane, SubCategory: "sub1", Category: constants.Metrics, SumValue: 25.0},
		}, nil)

		originalGetCategoryShareProportion := config.GetCategoryShareProportion
		getCategoryShareProportion = func(category constants.Category, cloudService constants.CloudService) float64 {
			return 0.5 // 50% share
		}
		defer func() { getCategoryShareProportion = originalGetCategoryShareProportion }()

		resp, err := ShareCostsStatisticReport(context.Background(), req)
		assert.NoError(t, err)
		assert.Len(t, resp.TenantStatisticReport, 1)
		// Cost = 100 * 0.5 * (25/50) = 100 * 0.5 * 0.5 = 25.0
		assert.Equal(t, 25.0, resp.TenantStatisticReport[0].Cost)
	})

	t.Run("ShareCostsStatisticReport GetMonthlyCost error", func(t *testing.T) {
		ctrl, _, _, mockStatisticReportRW, _, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		mockStatisticReportRW.EXPECT().GetMonthlyCost(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(0.0, errors.New("get monthly cost error"))

		_, err := ShareCostsStatisticReport(context.Background(), req)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "get monthly cost error")
	})

	t.Run("ShareCostsStatisticReport MonthlyStatisticSum error", func(t *testing.T) {
		ctrl, _, mockStatisticMergeRW, mockStatisticReportRW, _, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		mockStatisticReportRW.EXPECT().GetMonthlyCost(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(100.0, nil)
		mockStatisticMergeRW.EXPECT().MonthlyStatisticSum(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(0.0, errors.New("sum error"))

		_, err := ShareCostsStatisticReport(context.Background(), req)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "sum error")
	})

	t.Run("ShareCostsStatisticReport MonthlyTenantStatistic error", func(t *testing.T) {
		ctrl, _, mockStatisticMergeRW, mockStatisticReportRW, _, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		mockStatisticReportRW.EXPECT().GetMonthlyCost(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(100.0, nil)
		mockStatisticMergeRW.EXPECT().MonthlyStatisticSum(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(50.0, nil)
		mockStatisticMergeRW.EXPECT().MonthlyTenantStatistic(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, errors.New("tenant statistic error"))

		_, err := ShareCostsStatisticReport(context.Background(), req)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "tenant statistic error")
	})
}

func TestEnsureLocalData(t *testing.T) {
	configs.Init(structs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})

	// Create a temporary directory for testing file operations
	tempDir := t.TempDir()
	originalDataPath := fmt.Sprintf("%s/%s", "data", constants.StatisticTaskDataPrefix)
	os.Setenv("O11Y_DATA_PATH", tempDir) // Override base path for testing

	// Reset O11Y_DATA_PATH after test
	defer func() {
		os.Setenv("O11Y_DATA_PATH", originalDataPath)
	}()

	taskName := "test_task"
	timestamp := time.Now().Unix()

	t.Run("File exists", func(t *testing.T) {
		ctrl, _, _, _, mockScheduleTaskDataSync, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		filePath := filepath.Join(tempDir, constants.StatisticTaskDataPrefix, taskName, fmt.Sprintf("test_task_%d.data", timestamp))
		EnsurePath(filepath.Dir(filePath)) // Ensure directory exists for file creation
		err := os.WriteFile(filePath, []byte("test data"), 0644)
		assert.NoError(t, err)

		mockScheduleTaskDataSync.EXPECT().LocalDataFilePath(gomock.Any(), taskName, timestamp).Return(filePath).AnyTimes()

		returnedFilePath, err := EnsureLocalData(context.Background(), taskName, timestamp)
		assert.NoError(t, err)
		assert.Equal(t, filePath, returnedFilePath)
	})

	t.Run("File does not exist, download succeeds", func(t *testing.T) {
		ctrl, _, _, _, mockScheduleTaskDataSync, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		filePath := filepath.Join(tempDir, constants.StatisticTaskDataPrefix, taskName, fmt.Sprintf("test_task_%d.data", timestamp+1))
		// Ensure the directory is clean for this test run
		os.RemoveAll(filepath.Dir(filePath))

		mockScheduleTaskDataSync.EXPECT().LocalDataFilePath(gomock.Any(), taskName, timestamp+1).Return(filePath)
		mockScheduleTaskDataSync.EXPECT().Download(gomock.Any(), taskName, timestamp+1).Return(filePath, nil)

		returnedFilePath, err := EnsureLocalData(context.Background(), taskName, timestamp+1)
		assert.NoError(t, err)
		assert.Equal(t, filePath, returnedFilePath)
	})

	t.Run("File does not exist, download fails", func(t *testing.T) {
		ctrl, _, _, _, mockScheduleTaskDataSync, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		filePath := filepath.Join(tempDir, constants.StatisticTaskDataPrefix, taskName, fmt.Sprintf("test_task_%d.data", timestamp+2))
		// Ensure the directory is clean for this test run
		os.RemoveAll(filepath.Dir(filePath))

		mockScheduleTaskDataSync.EXPECT().LocalDataFilePath(gomock.Any(), taskName, timestamp+2).Return(filePath)
		mockScheduleTaskDataSync.EXPECT().Download(gomock.Any(), taskName, timestamp+2).Return("", errors.New("download failed"))

		returnedFilePath, err := EnsureLocalData(context.Background(), taskName, timestamp+2)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "download failed")
		assert.Empty(t, returnedFilePath)
	})
}

func TestCalculateMonthlyProportionSum(t *testing.T) {
	configs.Init(structs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})

	t.Run("CalculateMonthlyProportionSum success", func(t *testing.T) {
		ctrl, _, _, mockStatisticReportRW, _, mockS3Client := setupStatisticMocks(t)
		defer ctrl.Finish()

		timestamp := time.Now().Unix()
		req := message.MonthlyProportionSumReq{Timestamp: int32(timestamp)}
		monthString := time.Unix(timestamp, 0).Format("2006-01")

		mockProportionSums := []models.MonthlyProportionSum{
			{Timestamp: timestamp, Month: monthString, Provider: "aws", Env: "prod", MeteringGroup: "dedicated", Tenant: "t1", Proportion: 0.5},
		}

		mockStatisticReportRW.EXPECT().CalculateMonthlyProportionSum(gomock.Any(), gomock.Any(), gomock.Any()).Return(mockProportionSums, nil)
		mockS3Client.EXPECT().Put(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil)

		resp, err := CalculateMonthlyProportionSum(context.Background(), req)
		assert.NoError(t, err)
		assert.Equal(t, req.Timestamp, resp.Timestamp)
		assert.Contains(t, resp.File, ".csv")
	})

	t.Run("CalculateMonthlyProportionSum query error", func(t *testing.T) {
		ctrl, _, _, mockStatisticReportRW, _, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		timestamp := time.Now().Unix()
		req := message.MonthlyProportionSumReq{Timestamp: int32(timestamp)}

		mockStatisticReportRW.EXPECT().CalculateMonthlyProportionSum(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, errors.New("db error"))

		_, err := CalculateMonthlyProportionSum(context.Background(), req)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "db error")
	})

	t.Run("CalculateMonthlyProportionSum upload CSV error", func(t *testing.T) {
		ctrl, _, _, mockStatisticReportRW, _, mockS3Client := setupStatisticMocks(t)
		defer ctrl.Finish()

		timestamp := time.Now().Unix()
		req := message.MonthlyProportionSumReq{Timestamp: int32(timestamp)}
		monthString := time.Unix(timestamp, 0).Format("2006-01")

		mockProportionSums := []models.MonthlyProportionSum{
			{Timestamp: timestamp, Month: monthString, Provider: "aws", Env: "prod", MeteringGroup: "dedicated", Tenant: "t1", Proportion: 0.5},
		}

		mockStatisticReportRW.EXPECT().CalculateMonthlyProportionSum(gomock.Any(), gomock.Any(), gomock.Any()).Return(mockProportionSums, nil)
		mockS3Client.EXPECT().Put(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(errors.New("s3 upload error"))

		_, err := CalculateMonthlyProportionSum(context.Background(), req)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "s3 upload error")
	})
}

func TestUploadCSVFromProportionSum(t *testing.T) {
	configs.Init(structs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})

	t.Run("UploadCSVFromProportionSum success", func(t *testing.T) {
		ctrl, _, _, _, _, mockS3Client := setupStatisticMocks(t)
		defer ctrl.Finish()

		records := []models.MonthlyProportionSum{
			{Timestamp: 1672531200, Month: "2023-01", Provider: "aws", Env: "prod", MeteringGroup: "dedicated", Tenant: "t1", Proportion: 0.5},
			{Timestamp: 1672531200, Month: "2023-01", Provider: "aws", Env: "prod", MeteringGroup: "serverless", Tenant: "t2", Proportion: 0.3},
		}
		monthString := "2023-01"

		mockS3Client.EXPECT().Put(gomock.Any(), "o11y-prod-shared-us-west-2", "metering/split/o11y_usage_proportion_2023-01.csv", gomock.Any()).Return(nil)

		filePath, err := UploadCSVFromProportionSum(context.Background(), records, monthString)
		assert.NoError(t, err)
		assert.Equal(t, "metering/split/o11y_usage_proportion_2023-01.csv", filePath)
	})

	t.Run("UploadCSVFromProportionSum S3 error", func(t *testing.T) {
		ctrl, _, _, _, _, mockS3Client := setupStatisticMocks(t)
		defer ctrl.Finish()

		records := []models.MonthlyProportionSum{
			{Timestamp: 1672531200, Month: "2023-01", Provider: "aws", Env: "prod", MeteringGroup: "dedicated", Tenant: "t1", Proportion: 0.5},
		}
		monthString := "2023-01"

		mockS3Client.EXPECT().Put(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(errors.New("s3 put error"))

		filePath, err := UploadCSVFromProportionSum(context.Background(), records, monthString)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "s3 put error")
		assert.Empty(t, filePath)
	})
}

// --- Tests for statisticservice.go ---

// Mock TaskHandler for GenerateStatisticTask and ExecuteStatisticTask
type MockTaskHandler struct {
	BuildFunc   func(ctx context.Context, period string) (models.StatisticTask, error)
	ExecuteFunc func(ctx context.Context, task models.StatisticTask) ([]*models.StatisticRecord, error)
}

func (m *MockTaskHandler) Build(ctx context.Context, period string) (models.StatisticTask, error) {
	if m.BuildFunc != nil {
		return m.BuildFunc(ctx, period)
	}
	return models.StatisticTask{}, nil
}

func (m *MockTaskHandler) Execute(ctx context.Context, task models.StatisticTask) ([]*models.StatisticRecord, error) {
	if m.ExecuteFunc != nil {
		return m.ExecuteFunc(ctx, task)
	}
	return nil, nil
}

func TestGenerateStatisticTask(t *testing.T) {
	configs.Init(structs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})

	t.Run("GenerateStatisticTask success", func(t *testing.T) {
		ctrl, mockStatisticRW, _, _, _, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		req := message.GenerateStatisticTaskReq{Period: "2023-01-01 00:00:00"}

		// Mock getTasks to return a specific task type
		originalGetTasks := getTasks
		getTasks = func() []config.TaskConfig {
			return []config.TaskConfig{
				{StatisticClass: constants.DryRunStatisticClass, MeteringGroup: constants.DryRunMeteringGroup, Env: constants.Prod},
			}
		}
		defer func() { getTasks = originalGetTasks }()

		// Mock the TaskHandler's Build method
		originalPickTaskHandler := pickTaskHandler
		pickHandler = func(statisticClass constants.StatisticClass, meteringGroup constants.MeteringGroup, bizEnv constants.BizEnv) handler.TaskHandler {
			return &MockTaskHandler{
				BuildFunc: func(ctx context.Context, period string) (models.StatisticTask, error) {
					return models.StatisticTask{ID: "task-1", Period: period, State: constants.StatePending}, nil
				},
			}
		}
		defer func() { pickHandler = originalPickTaskHandler }()

		mockStatisticRW.EXPECT().BatchCreateTask(gomock.Any(), gomock.Any()).DoAndReturn(
			func(_ context.Context, tasks []*models.StatisticTask) ([]*models.StatisticTask, error) {
				for i, task := range tasks {
					task.ID = fmt.Sprintf("task-%d", i+1) // Assign IDs as if created in DB
				}
				return tasks, nil
			},
		)

		resp, err := GenerateStatisticTask(context.Background(), req)
		assert.NoError(t, err)
		assert.Len(t, resp.TaskIDs, 1)
		assert.Equal(t, "task-1", resp.TaskIDs[0])
	})

	t.Run("GenerateStatisticTask build error", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Region: "us-east-1",
			Vendor: "aws",
		})
		req := message.GenerateStatisticTaskReq{Period: "2023-01-01 00:00:00"}

		originalGetTasks := getTasks
		getTasks = func() []config.TaskConfig {
			return []config.TaskConfig{
				{StatisticClass: constants.DryRunStatisticClass, MeteringGroup: constants.DryRunMeteringGroup, Env: constants.Prod},
			}
		}
		defer func() { getTasks = originalGetTasks }()

		originalPickTaskHandler := pickTaskHandler
		pickHandler = func(statisticClass constants.StatisticClass, meteringGroup constants.MeteringGroup, bizEnv constants.BizEnv) handler.TaskHandler {
			return &MockTaskHandler{
				BuildFunc: func(ctx context.Context, period string) (models.StatisticTask, error) {
					return models.StatisticTask{}, errors.New("build error")
				},
			}
		}
		defer func() { pickHandler = originalPickTaskHandler }()

		resp, err := GenerateStatisticTask(context.Background(), req)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "build error")
		assert.Equal(t, message.GenerateStatisticTaskResp{}, resp)
	})

	t.Run("GenerateStatisticTask BatchCreateTask error", func(t *testing.T) {
		ctrl, mockStatisticRW, _, _, _, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		req := message.GenerateStatisticTaskReq{Period: "2023-01-01 00:00:00"}

		originalGetTasks := getTasks
		getTasks = func() []config.TaskConfig {
			return []config.TaskConfig{
				{StatisticClass: constants.DryRunStatisticClass, MeteringGroup: constants.DryRunMeteringGroup, Env: constants.Prod},
			}
		}
		defer func() { getTasks = originalGetTasks }()

		originalPickTaskHandler := pickTaskHandler
		pickHandler = func(statisticClass constants.StatisticClass, meteringGroup constants.MeteringGroup, bizEnv constants.BizEnv) handler.TaskHandler {
			return &MockTaskHandler{
				BuildFunc: func(ctx context.Context, period string) (models.StatisticTask, error) {
					return models.StatisticTask{ID: "task-1", Period: period, State: constants.StatePending}, nil
				},
			}
		}
		defer func() { pickHandler = originalPickTaskHandler }()

		mockStatisticRW.EXPECT().BatchCreateTask(gomock.Any(), gomock.Any()).Return(nil, errors.New("batch create error"))

		_, err := GenerateStatisticTask(context.Background(), req)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "batch create error")
	})
}

func TestExecuteStatisticTask(t *testing.T) {
	configs.Init(structs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})

	t.Run("ExecuteStatisticTask success", func(t *testing.T) {
		ctrl, mockStatisticRW, _, _, _, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		req := message.ExecuteStatisticTaskReq{ID: "test-task-id"}
		mockTask := &models.StatisticTask{
			ID: "test-task-id", StatisticClass: constants.DryRunStatisticClass,
			MeteringGroup: constants.DryRunMeteringGroup, Env: constants.Prod,
			Period: "2023-01-01 00:00:00", State: constants.StatePending,
		}
		mockRecords := []*models.StatisticRecord{
			{TaskID: "test-task-id", Value: 1.0},
		}

		mockStatisticRW.EXPECT().GetTask(gomock.Any(), req.ID).Return(mockTask, nil)
		mockStatisticRW.EXPECT().UpdateTaskState(gomock.Any(), mockTask.ID, constants.StatePending, constants.StateProcessing).Return(nil)

		originalPickTaskHandler := pickTaskHandler
		pickHandler = func(statisticClass constants.StatisticClass, meteringGroup constants.MeteringGroup, bizEnv constants.BizEnv) handler.TaskHandler {
			return &MockTaskHandler{
				ExecuteFunc: func(ctx context.Context, task models.StatisticTask) ([]*models.StatisticRecord, error) {
					return mockRecords, nil
				},
			}
		}
		defer func() { pickHandler = originalPickTaskHandler }()

		mockStatisticRW.EXPECT().SaveStatisticRecordsForTask(gomock.Any(), mockTask.ID, mockRecords).Return(mockRecords, nil)

		resp, err := ExecuteStatisticTask(context.Background(), req)
		assert.NoError(t, err)
		assert.Equal(t, int32(1), resp.RecordsCount)
	})

	t.Run("ExecuteStatisticTask GetTask error", func(t *testing.T) {
		ctrl, mockStatisticRW, _, _, _, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		req := message.ExecuteStatisticTaskReq{ID: "test-task-id"}
		mockStatisticRW.EXPECT().GetTask(gomock.Any(), req.ID).Return(nil, errors.New("get task error"))

		_, err := ExecuteStatisticTask(context.Background(), req)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "get task error")
	})

	t.Run("ExecuteStatisticTask UpdateTaskState to Processing error", func(t *testing.T) {
		ctrl, mockStatisticRW, _, _, _, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		req := message.ExecuteStatisticTaskReq{ID: "test-task-id"}
		mockTask := &models.StatisticTask{ID: "test-task-id", State: constants.StatePending}
		mockStatisticRW.EXPECT().GetTask(gomock.Any(), req.ID).Return(mockTask, nil)
		mockStatisticRW.EXPECT().UpdateTaskState(gomock.Any(), mockTask.ID, constants.StatePending, constants.StateProcessing).Return(errors.New("update state error"))

		_, err := ExecuteStatisticTask(context.Background(), req)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "update state error")
	})

	t.Run("ExecuteStatisticTask handler Execute error", func(t *testing.T) {
		ctrl, mockStatisticRW, _, _, _, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		req := message.ExecuteStatisticTaskReq{ID: "test-task-id"}
		mockTask := &models.StatisticTask{
			ID: "test-task-id", StatisticClass: constants.DryRunStatisticClass,
			MeteringGroup: constants.DryRunMeteringGroup, Env: constants.Prod,
			Period: "2023-01-01 00:00:00", State: constants.StatePending,
		}
		mockStatisticRW.EXPECT().GetTask(gomock.Any(), req.ID).Return(mockTask, nil)
		mockStatisticRW.EXPECT().UpdateTaskState(gomock.Any(), mockTask.ID, constants.StatePending, constants.StateProcessing).Return(nil)

		originalPickTaskHandler := pickTaskHandler
		pickHandler = func(statisticClass constants.StatisticClass, meteringGroup constants.MeteringGroup, bizEnv constants.BizEnv) handler.TaskHandler {
			return &MockTaskHandler{
				ExecuteFunc: func(ctx context.Context, task models.StatisticTask) ([]*models.StatisticRecord, error) {
					return nil, errors.New("execute handler error")
				},
			}
		}
		defer func() { pickHandler = originalPickTaskHandler }()

		mockStatisticRW.EXPECT().UpdateTaskState(gomock.Any(), mockTask.ID, constants.StateProcessing, constants.StateFailed).Return(nil)

		_, err := ExecuteStatisticTask(context.Background(), req)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "execute handler error")
	})

	t.Run("ExecuteStatisticTask SaveStatisticRecordsForTask error", func(t *testing.T) {
		ctrl, mockStatisticRW, _, _, _, _ := setupStatisticMocks(t)
		defer ctrl.Finish()

		req := message.ExecuteStatisticTaskReq{ID: "test-task-id"}
		mockTask := &models.StatisticTask{
			ID: "test-task-id", StatisticClass: constants.DryRunStatisticClass,
			MeteringGroup: constants.DryRunMeteringGroup, Env: constants.Prod,
			Period: "2023-01-01 00:00:00", State: constants.StatePending,
		}
		mockRecords := []*models.StatisticRecord{
			{TaskID: "test-task-id", Value: 1.0},
		}

		mockStatisticRW.EXPECT().GetTask(gomock.Any(), req.ID).Return(mockTask, nil)
		mockStatisticRW.EXPECT().UpdateTaskState(gomock.Any(), mockTask.ID, constants.StatePending, constants.StateProcessing).Return(nil)

		originalPickTaskHandler := pickTaskHandler
		pickHandler = func(statisticClass constants.StatisticClass, meteringGroup constants.MeteringGroup, bizEnv constants.BizEnv) handler.TaskHandler {
			return &MockTaskHandler{
				ExecuteFunc: func(ctx context.Context, task models.StatisticTask) ([]*models.StatisticRecord, error) {
					return mockRecords, nil
				},
			}
		}
		defer func() { pickHandler = originalPickTaskHandler }()

		mockStatisticRW.EXPECT().SaveStatisticRecordsForTask(gomock.Any(), mockTask.ID, mockRecords).Return(nil, errors.New("save records error"))
		// Note: UpdateTaskState to StateFailed is not called in this specific error path as per the original code.

		_, err := ExecuteStatisticTask(context.Background(), req)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "save records error")
	})
}