package service

import (
	"context"
	"os"

	"github.com/stretchr/testify/mock"
	"github.com/tidbcloud/pov-service/internal/common/params"
	"github.com/tidbcloud/pov-service/internal/model"
	"github.com/tidbcloud/pov-service/pkg/browser"
)

// MockWorkloadService is a mock implementation of service.Workload
type MockWorkloadService struct {
	mock.Mock
}

func (m *MockWorkloadService) ListWorkloads(ctx context.Context, param params.ListWorkloads) ([]model.Workload, error) {
	args := m.Called(ctx)
	return args.Get(0).([]model.Workload), args.Error(1)
}

func (m *MockWorkloadService) GetWorkload(ctx context.Context, id int64) (model.Workload, error) {
	args := m.Called(ctx)
	return args.Get(0).(model.Workload), args.Error(1)
}

func (m *MockWorkloadService) RunWorkloadTest(ctx context.Context, param params.RunWorkload, startedBy string) (model.WorkloadTest, error) {
	args := m.Called(ctx, param, startedBy)
	return args.Get(0).(model.WorkloadTest), args.Error(1)
}

func (m *MockWorkloadService) GetWorkloadTest(ctx context.Context, ID string) (model.WorkloadTest, error) {
	args := m.Called(ctx, ID)
	return args.Get(0).(model.WorkloadTest), args.Error(1)
}

func (m *MockWorkloadService) CreateWorkloadTest(ctx context.Context, param params.RunWorkload, startedBy string) (model.WorkloadTest, error) {
	args := m.Called(ctx, param, startedBy)
	return args.Get(0).(model.WorkloadTest), args.Error(1)
}

func (m *MockWorkloadService) TriggerCreateCluster(ctx context.Context, ID string) (model.WorkloadTest, error) {
	args := m.Called(ctx, ID)
	ret, err := args.Get(0).(model.WorkloadTest), args.Error(1)
	ret.OnStart()

	return ret, err
}

func (m *MockWorkloadService) TriggerPrepareData(ctx context.Context, ID string) (model.WorkloadTest, error) {
	args := m.Called(ctx, ID)
	return args.Get(0).(model.WorkloadTest), args.Error(1)
}

func (m *MockWorkloadService) TriggerRunTraffic(ctx context.Context, ID string) (model.WorkloadTest, error) {
	args := m.Called(ctx, ID)
	return args.Get(0).(model.WorkloadTest), args.Error(1)
}

func (m *MockWorkloadService) ListWorkloadTests(ctx context.Context, param params.ListWorkloadTests) ([]model.WorkloadTest, error) {
	args := m.Called(ctx, param)
	return args.Get(0).([]model.WorkloadTest), args.Error(1)
}

func (m *MockWorkloadService) StopWorkloadTest(ctx context.Context, ID string) (model.WorkloadTest, error) {
	args := m.Called(ctx, ID)
	return args.Get(0).(model.WorkloadTest), args.Error(1)
}

func (m *MockWorkloadService) ScaleCluster(ctx context.Context, ID string, replicas int) error {
	args := m.Called(ctx, ID, replicas)
	return args.Error(0)
}

func (m *MockWorkloadService) UpdateClusterRU(ctx context.Context, testID string, minRU, maxRU int64) error {
	args := m.Called(ctx, testID, minRU, maxRU)
	return args.Error(0)
}

func (m *MockWorkloadService) UpdateTrafficMaxConn(
	ctx context.Context,
	testID string,
	maxConNumber int,
) error {
	args := m.Called(ctx, testID, maxConNumber)
	return args.Error(0)
}

func (m *MockWorkloadService) UpdateWorkloadTest(
	ctx context.Context,
	testID string,
	params params.UpdateWorkloadTest,
) error {
	args := m.Called(ctx, testID, params)
	return args.Error(0)
}

func (m *MockWorkloadService) InstructWorkloadTest(
	ctx context.Context,
	testID string,
	params params.InstructWorkloadTest,
) error {
	args := m.Called(ctx, testID, params)
	return args.Error(0)
}

func (m *MockWorkloadService) SampleDBs(
	ctx context.Context,
	testID string,
	options model.SampleDBOptions) (model.DBSampleResult, error) {
	args := m.Called(ctx, testID, options)
	return args.Get(0).(model.DBSampleResult), args.Error(1)
}

func (m *MockWorkloadService) SampleData(
	ctx context.Context,
	testID string,
	options model.SampleDataOptions) (model.TableQueryResult, error) {
	args := m.Called(ctx, testID, options)
	return args.Get(0).(model.TableQueryResult), args.Error(1)
}

func (m *MockWorkloadService) ListAuditLog(
	ctx context.Context,
	testID string,
) ([]model.AuditLog, error) {
	args := m.Called(ctx, testID)
	return args.Get(0).([]model.AuditLog), args.Error(1)
}

func (m *MockWorkloadService) PreserveWorkloadSample(
	ctx context.Context,
	workloadID int64,
	cluster model.TiDBCluster) error {
	args := m.Called(ctx, workloadID, cluster)

	return args.Error(0)
}

func (m *MockWorkloadService) GetPreservedSampleData(
	ctx context.Context,
	workloadID int64,
	dbName, tableName string,
) (model.TableQueryResult, error) {
	args := m.Called(ctx, workloadID, dbName, tableName)
	return args.Get(0).(model.TableQueryResult), args.Error(1)
}

func (m *MockWorkloadService) GetPreservedSampleDBs(
	ctx context.Context,
	workloadID int64) (model.DBSampleResult, error) {
	args := m.Called(ctx, workloadID)
	return args.Get(0).(model.DBSampleResult), args.Error(1)
}

func (m *MockWorkloadService) DownloadTestPDF(
	ctx context.Context,
	testID string,
	opts browser.PrintPDFOptions,
) (*os.File, error) {
	args := m.Called(ctx, testID, opts)
	return args.Get(0).(*os.File), args.Error(1)
}

func (m *MockWorkloadService) GetWorkloadUsage(
	ctx context.Context) (map[int64]int, error) {
	args := m.Called(ctx)
	return args.Get(0).(map[int64]int), args.Error(1)
}

func (m *MockWorkloadService) InjectFailure(
	ctx context.Context, testID string, duration string) error {
	args := m.Called(ctx, testID, duration)
	return args.Error(0)

}

func (m *MockWorkloadService) RollbackFailure(ctx context.Context, testID string) error {
	args := m.Called(ctx, testID)
	return args.Error(0)
}