package mysql

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/tidbcloud/pov-service/internal/common/options"
	"github.com/tidbcloud/pov-service/internal/model"
	"github.com/tidbcloud/pov-service/internal/pkg/dbmarshal"
	"github.com/tidbcloud/pov-service/internal/repo"
)

type workloadTestRepo struct {
	db *sqlx.DB
	TxnProvider
}

// WorkloadTestDB represents the database structure for workload_tests table
type WorkloadTestDB struct {
	ID           string `db:"id"`
	Name         string `db:"name"`
	WorkloadID   int64  `db:"workload_id"`
	WorkloadName string `db:"workload_name"`
	ParamsJSON   string `db:"params"`

	// ExecBase fields
	CreatedAt  time.Time  `db:"created_at"`
	UpdatedAt  *time.Time `db:"updated_at"`
	StartedAt  *time.Time `db:"started_at"`
	FinishedAt *time.Time `db:"finished_at"`
	Status     string     `db:"status"`
	Message    string     `db:"message"`

	// WorkloadTest specific fields
	Progress  float64        `db:"progress"`
	Phase     string         `db:"phase"`
	IsAuto    bool           `db:"is_auto"`
	StartedBy sql.NullString `db:"started_by"`

	// Cluster information as JSON
	ClusterJSON sql.NullString `db:"cluster"`
}

// WorkloadTestTaskDB represents the database structure for workload_test_tasks table
type WorkloadTestTaskDB struct {
	ID               int64          `db:"id"`
	WorkloadTestID   string         `db:"workload_test_id"`
	TaskType         string         `db:"task_type"`
	CreatedAt        time.Time      `db:"created_at"`
	UpdatedAt        *time.Time     `db:"updated_at"`
	StartedAt        *time.Time     `db:"started_at"`
	FinishedAt       *time.Time     `db:"finished_at"`
	Status           string         `db:"status"`
	Message          string         `db:"message"`
	Progress         float64        `db:"progress"`
	RelatedServiceID string         `db:"related_service_id"`
	Phase            string         `db:"phase"`
	ResultJSON       sql.NullString `db:"result"`
}

// WorkloadTestParamDB represents the database structure for workload_test_params table
type WorkloadTestParamDB struct {
	ID             int64  `db:"id"`
	WorkloadTestID string `db:"workload_test_id"`
	ParamName      string `db:"param_name"`
	ParamValue     string `db:"param_value"`
}

const (
	workloadTestColumns = `
		id, name, workload_id, workload_name, params, status, message, progress, phase, is_auto, started_by, cluster, created_at, updated_at, started_at, finished_at
	`
	listWorkloadTestColumns = `
		id, name, workload_id, workload_name, status, message, progress, phase, is_auto, started_at, finished_at, started_by, cluster
	`

	workloadTestTasksColumns = `
		id, workload_test_id, task_type, status, message, progress, related_service_id, phase, result, created_at, updated_at, started_at, finished_at
	`
	workloadTestParamsColumns = `
		id, workload_test_id, param_name, param_value
	`
)

// NewWorkloadTest creates a new MySQL-based WorkloadTest repository
func NewWorkloadTest(db *sqlx.DB) repo.WorkloadTest {
	return &workloadTestRepo{
		db: db,
		TxnProvider: TxnProvider{
			db: db,
		},
	}
}

func (r *workloadTestRepo) Create(ctx context.Context, workload model.WorkloadTest) (model.WorkloadTest, error) {
	tx, err := r.db.BeginTxx(ctx, nil)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	// Check if workload test already exists
	var exists bool
	err = tx.GetContext(ctx, &exists, "SELECT EXISTS(SELECT 1 FROM workload_tests WHERE id = ?)", workload.ID)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to check if workload test exists: %w", err)
	}
	if exists {
		return model.WorkloadTest{}, fmt.Errorf("workload test %s already exists", workload.ID)
	}

	// NOTE: created_at and updated_at are now handled by database defaults
	// Only started_at and finished_at are controlled by the service

	// Serialize params to JSON
	paramsJSON, err := json.Marshal(workload.Params)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to marshal params: %w", err)
	}

	// Serialize cluster to JSON using custom marshaling if it exists
	var clusterJSON sql.NullString
	if workload.Cluster != nil {
		clusterBytes, err := dbmarshal.Marshal(workload.Cluster)
		if err != nil {
			return model.WorkloadTest{}, fmt.Errorf("failed to marshal cluster: %w", err)
		}
		clusterJSON = sql.NullString{String: string(clusterBytes), Valid: true}
	}

	// Insert main workload test record - let database handle created_at and updated_at
	query := `
		INSERT INTO workload_tests (
			id, name, workload_id, workload_name, params,
			started_at, finished_at, status, message,
			progress, phase, is_auto, started_by, cluster
		) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`

	_, err = tx.ExecContext(ctx, query,
		workload.ID, workload.Name, workload.WorkloadID, workload.WorkloadName, string(paramsJSON),
		workload.StartedAt, workload.FinishedAt,
		workload.Status, workload.Message,
		workload.Progress, workload.Phase, workload.IsAuto, workload.StartedBy, clusterJSON)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to insert workload test: %w", err)
	}

	// Insert tasks if any
	if err := r.insertTasks(ctx, tx, workload.ID, workload.Tasks); err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to insert tasks: %w", err)
	}

	// Fetch the record back to get database-generated timestamps
	var dbWorkload WorkloadTestDB
	getQuery := fmt.Sprintf("SELECT %s FROM workload_tests WHERE id = ?", workloadTestColumns)
	err = tx.GetContext(ctx, &dbWorkload, getQuery, workload.ID)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to fetch created workload: %w", err)
	}

	// Get tasks within the transaction before committing
	tasks, err := r.getTasksInTx(ctx, tx, workload.ID)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to get tasks after create: %w", err)
	}

	if err = tx.Commit(); err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to commit transaction: %w", err)
	}

	// Convert back to model with database timestamps
	result, err := r.dbToModel(dbWorkload)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to convert db record to model: %w", err)
	}

	result.Tasks = tasks
	return result, nil
}

func (r *workloadTestRepo) Update(
	ctx context.Context,
	workload model.WorkloadTest,
	opts ...options.TableStatement,
) (model.WorkloadTest, error) {
	tx, err := r.db.BeginTxx(ctx, nil)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	tblOptions := options.NewTableOption(opts...)
	waitCmd := "NOWAIT"
	if tblOptions.ShouldWait {
		waitCmd = ""
	}

	// Check if workload exists and lock it to prevent concurrent updates
	var currentVersion WorkloadTestDB
	lockQuery := fmt.Sprintf("SELECT %s FROM workload_tests WHERE id = ? FOR UPDATE %s", workloadTestColumns, waitCmd)
	err = tx.GetContext(ctx, &currentVersion, lockQuery, workload.ID)
	if err != nil {
		if err == sql.ErrNoRows {
			return model.WorkloadTest{}, fmt.Errorf("workload %s not found", workload.ID)
		}
		return model.WorkloadTest{}, fmt.Errorf("failed to lock workload for update: %w", err)
	}

	// NOTE: updated_at is now handled by database ON UPDATE CURRENT_TIMESTAMP
	// Only started_at and finished_at are controlled by the service

	// Serialize params to JSON
	paramsJSON, err := json.Marshal(workload.Params)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to marshal params: %w", err)
	}

	// Serialize cluster to JSON using custom marshaling if it exists
	var clusterJSON sql.NullString
	if workload.Cluster != nil {
		clusterBytes, err := dbmarshal.Marshal(workload.Cluster)
		if err != nil {
			return model.WorkloadTest{}, fmt.Errorf("failed to marshal cluster: %w", err)
		}
		clusterJSON = sql.NullString{String: string(clusterBytes), Valid: true}
	}

	// Update main workload test record - let database handle updated_at
	// Note: started_by is not updated as it should be immutable after creation
	query := `
		UPDATE workload_tests SET
			name = ?, workload_id = ?, workload_name = ?, params = ?,
			started_at = ?, finished_at = ?, status = ?, message = ?,
			progress = ?, phase = ?, is_auto = ?, cluster = ?
		WHERE id = ?`

	_, err = tx.ExecContext(ctx, query,
		workload.Name, workload.WorkloadID, workload.WorkloadName, string(paramsJSON),
		workload.StartedAt, workload.FinishedAt, workload.Status, workload.Message,
		workload.Progress, workload.Phase, workload.IsAuto, clusterJSON,
		workload.ID)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to update workload test: %w", err)
	}

	// Upsert tasks (update existing, insert new)
	if err := r.upsertTasks(ctx, tx, workload.ID, workload.Tasks); err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to upsert tasks: %w", err)
	}

	// Fetch the record back to get database-generated updated_at timestamp
	var dbWorkload WorkloadTestDB
	getQuery := fmt.Sprintf("SELECT %s FROM workload_tests WHERE id = ?", workloadTestColumns)
	err = tx.GetContext(ctx, &dbWorkload, getQuery, workload.ID)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to fetch updated workload: %w", err)
	}

	// Get tasks within the transaction before committing
	tasks, err := r.getTasksInTx(ctx, tx, workload.ID)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to get tasks after update: %w", err)
	}

	if err = tx.Commit(); err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to commit transaction: %w", err)
	}

	// Convert back to model with database timestamps
	result, err := r.dbToModel(dbWorkload)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to convert db record to model: %w", err)
	}

	result.Tasks = tasks
	return result, nil
}

func (r *workloadTestRepo) Get(ctx context.Context, ID string) (model.WorkloadTest, error) {
	tx, err := r.db.BeginTxx(ctx, nil)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	// Get main workload test record
	var dbWorkload WorkloadTestDB
	query := fmt.Sprintf("SELECT %s FROM workload_tests WHERE id = ?", workloadTestColumns)
	err = tx.GetContext(ctx, &dbWorkload, query, ID)
	if err != nil {
		if err == sql.ErrNoRows {
			return model.WorkloadTest{}, fmt.Errorf("workload %s not found", ID)
		}
		return model.WorkloadTest{}, fmt.Errorf("failed to get workload test: %w", err)
	}

	// Convert database record to model
	workload, err := r.dbToModel(dbWorkload)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to convert db record to model: %w", err)
	}

	// Get tasks
	tasks, err := r.getTasksInTx(ctx, tx, ID)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to get tasks: %w", err)
	}
	workload.Tasks = tasks

	if err = tx.Commit(); err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to commit transaction: %w", err)
	}

	return workload, nil
}

// NOTE: List will only return workload tests with subset of all fields by default. see @listWorkloadTestColumns
func (r *workloadTestRepo) List(
	ctx context.Context,
	options options.ListOption[options.WorkloadTestRefOption]) ([]model.WorkloadTest, error) {
	tx, err := r.db.BeginTxx(ctx, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	condition, args := buildQueries(options.TableOption)
	selectStr := listWorkloadTestColumns
	if len(options.Selects) > 0 {
		selectStr = options.SelectString()
	}
	query := fmt.Sprintf(
		"SELECT %s FROM workload_tests %s", selectStr, condition)

	var dbWorkloads []WorkloadTestDB
	err = tx.SelectContext(ctx, &dbWorkloads, query, args...)
	if err != nil {
		return nil, fmt.Errorf("failed to list active workloads: %w", err)
	}

	var workloads []model.WorkloadTest
	for _, dbWorkload := range dbWorkloads {
		workload, err := r.dbToModel(dbWorkload)
		if err != nil {
			return nil, fmt.Errorf("failed to convert db record to model: %w", err)
		}

		if options.RefOption.WithTasks {
			// Get tasks for each workload
			tasks, err := r.getTasksInTx(ctx, tx, workload.ID)
			if err != nil {
				return nil, fmt.Errorf("failed to get tasks for workload %s: %w", workload.ID, err)
			}
			workload.Tasks = tasks
		}

		workloads = append(workloads, workload)
	}

	if err = tx.Commit(); err != nil {
		return nil, fmt.Errorf("failed to commit transaction: %w", err)
	}

	return workloads, nil
}

func (r *workloadTestRepo) ListActive(ctx context.Context) ([]model.WorkloadTest, error) {
	tx, err := r.db.BeginTxx(ctx, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	query := fmt.Sprintf("SELECT %s FROM workload_tests WHERE status IN ('pending', 'running', 'stopping') ORDER BY created_at DESC", workloadTestColumns)

	var dbWorkloads []WorkloadTestDB
	err = tx.SelectContext(ctx, &dbWorkloads, query)
	if err != nil {
		return nil, fmt.Errorf("failed to list active workloads: %w", err)
	}

	var workloads []model.WorkloadTest
	for _, dbWorkload := range dbWorkloads {
		workload, err := r.dbToModel(dbWorkload)
		if err != nil {
			return nil, fmt.Errorf("failed to convert db record to model: %w", err)
		}

		// Get tasks for each workload
		tasks, err := r.getTasksInTx(ctx, tx, workload.ID)
		if err != nil {
			return nil, fmt.Errorf("failed to get tasks for workload %s: %w", workload.ID, err)
		}
		workload.Tasks = tasks
		workloads = append(workloads, workload)
	}

	if err = tx.Commit(); err != nil {
		return nil, fmt.Errorf("failed to commit transaction: %w", err)
	}

	return workloads, nil
}

// Helper methods
func (r *workloadTestRepo) insertTasks(ctx context.Context, tx *sqlx.Tx, workloadTestID string, tasks []model.WorkloadTask[any]) error {
	if len(tasks) == 0 {
		return nil
	}

	// Let database handle created_at and updated_at for tasks too
	query := `
			INSERT INTO workload_test_tasks (
				workload_test_id, task_type,
				started_at, finished_at, status, message,
				progress, related_service_id, phase, result
			) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`

	for _, task := range tasks {
		var resultJSON sql.NullString
		if task.Result != nil {
			resultBytes, err := dbmarshal.Marshal(task.Result)
			if err != nil {
				return fmt.Errorf("failed to marshal task result: %w", err)
			}
			resultJSON = sql.NullString{String: string(resultBytes), Valid: true}
		}

		_, err := tx.ExecContext(ctx, query,
			workloadTestID, task.Typ,
			task.StartedAt, task.FinishedAt,
			task.Status, task.Message,
			task.Progress, task.RelatedServiceID, task.Phase, resultJSON)
		if err != nil {
			return fmt.Errorf("failed to insert task: %w", err)
		}
	}

	return nil
}

func (r *workloadTestRepo) upsertTasks(ctx context.Context, tx *sqlx.Tx, workloadTestID string, tasks []model.WorkloadTask[any]) error {
	if len(tasks) == 0 {
		return nil
	}

	for _, task := range tasks {
		var resultJSON sql.NullString
		if task.Result != nil {
			resultBytes, err := dbmarshal.Marshal(task.Result)
			if err != nil {
				return fmt.Errorf("failed to marshal task result: %w", err)
			}
			resultJSON = sql.NullString{String: string(resultBytes), Valid: true}
		}

		if task.ID > 0 {
			// Update existing task
			updateQuery := `
				UPDATE workload_test_tasks SET
					task_type = ?, started_at = ?, finished_at = ?, status = ?, message = ?,
					progress = ?, related_service_id = ?, phase = ?, result = ?
				WHERE id = ? AND workload_test_id = ?`

			_, err := tx.ExecContext(ctx, updateQuery,
				task.Typ, task.StartedAt, task.FinishedAt,
				task.Status, task.Message, task.Progress,
				task.RelatedServiceID, task.Phase, resultJSON,
				task.ID, workloadTestID)
			if err != nil {
				return fmt.Errorf("failed to update task %d: %w", task.ID, err)
			}
		} else {
			// Insert new task
			insertQuery := `
				INSERT INTO workload_test_tasks (
					workload_test_id, task_type,
					started_at, finished_at, status, message,
					progress, related_service_id, phase, result
				) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`

			_, err := tx.ExecContext(ctx, insertQuery,
				workloadTestID, task.Typ,
				task.StartedAt, task.FinishedAt,
				task.Status, task.Message,
				task.Progress, task.RelatedServiceID, task.Phase, resultJSON)
			if err != nil {
				return fmt.Errorf("failed to insert new task: %w", err)
			}
		}
	}

	return nil
}

func (r *workloadTestRepo) getTasksInTx(ctx context.Context, tx *sqlx.Tx, workloadTestID string) ([]model.WorkloadTask[any], error) {
	query := fmt.Sprintf("SELECT %s FROM workload_test_tasks WHERE workload_test_id = ? ORDER BY id", workloadTestTasksColumns)

	var dbTasks []WorkloadTestTaskDB
	err := tx.SelectContext(ctx, &dbTasks, query, workloadTestID)
	if err != nil {
		return nil, fmt.Errorf("failed to get tasks: %w", err)
	}

	return r.dbTasksToModel(dbTasks)
}

func (r *workloadTestRepo) dbTasksToModel(dbTasks []WorkloadTestTaskDB) ([]model.WorkloadTask[any], error) {
	var tasks []model.WorkloadTask[any]
	for _, dbTask := range dbTasks {
		task := model.WorkloadTask[any]{
			WorkloadTaskBase: model.WorkloadTaskBase{
				ID:  dbTask.ID,
				Typ: model.WorkloadTaskType(dbTask.TaskType),
				ExecBase: model.ExecBase{
					CreatedAt:  dbTask.CreatedAt,
					UpdatedAt:  dbTask.UpdatedAt,
					StartedAt:  dbTask.StartedAt,
					FinishedAt: dbTask.FinishedAt,
					Status:     model.ExecStatus(dbTask.Status),
					Message:    dbTask.Message,
				},
				Progress:         dbTask.Progress,
				Phase:            model.WorkloadTaskPhase(dbTask.Phase),
				RelatedServiceID: dbTask.RelatedServiceID,
			},
		}

		// Deserialize result based on task type
		if dbTask.ResultJSON.Valid {
			switch model.WorkloadTaskType(dbTask.TaskType) {
			case model.WorkloadTaskTypeCreateCluster:
				var result model.CreateClusterResult
				if err := dbmarshal.Unmarshal([]byte(dbTask.ResultJSON.String), &result); err != nil {
					return nil, fmt.Errorf("failed to unmarshal create cluster result: %w", err)
				}
				task.Result = result
			case model.WorkloadTaskTypePrepareData:
				var result model.PrepareDataResult
				if err := dbmarshal.Unmarshal([]byte(dbTask.ResultJSON.String), &result); err != nil {
					return nil, fmt.Errorf("failed to unmarshal prepare data result: %w", err)
				}
				task.Result = result
			case model.WorkloadTaskTypeRunTraffic:
				var result model.RunTrafficResult
				if err := dbmarshal.Unmarshal([]byte(dbTask.ResultJSON.String), &result); err != nil {
					return nil, fmt.Errorf("failed to unmarshal run traffic result: %w", err)
				}
				task.Result = result
			default:
				// For unknown task types, store as raw JSON
				var result any
				if err := dbmarshal.Unmarshal([]byte(dbTask.ResultJSON.String), &result); err != nil {
					return nil, fmt.Errorf("failed to unmarshal generic result: %w", err)
				}
				task.Result = result
			}
		}

		tasks = append(tasks, task)
	}

	return tasks, nil
}

func (r *workloadTestRepo) dbToModel(dbWorkload WorkloadTestDB) (model.WorkloadTest, error) {
	// Handle nullable StartedBy field
	var startedBy string
	if dbWorkload.StartedBy.Valid {
		startedBy = dbWorkload.StartedBy.String
	}

	workload := model.WorkloadTest{
		ID:           dbWorkload.ID,
		Name:         dbWorkload.Name,
		WorkloadID:   dbWorkload.WorkloadID,
		WorkloadName: dbWorkload.WorkloadName,
		ExecBase: model.ExecBase{
			CreatedAt:  dbWorkload.CreatedAt,
			UpdatedAt:  dbWorkload.UpdatedAt,
			StartedAt:  dbWorkload.StartedAt,
			FinishedAt: dbWorkload.FinishedAt,
			Status:     model.ExecStatus(dbWorkload.Status),
			Message:    dbWorkload.Message,
		},
		Progress:  dbWorkload.Progress,
		Phase:     model.WorkloadTestPhase(dbWorkload.Phase),
		IsAuto:    dbWorkload.IsAuto,
		StartedBy: startedBy,
	}

	// Deserialize params from JSON
	if dbWorkload.ParamsJSON != "" {
		var params model.WorkloadParams
		if err := json.Unmarshal([]byte(dbWorkload.ParamsJSON), &params); err != nil {
			return model.WorkloadTest{}, fmt.Errorf("failed to unmarshal params: %w", err)
		}
		workload.Params = params
	}

	// Deserialize cluster from JSON
	if dbWorkload.ClusterJSON.Valid {
		var cluster model.TiDBCluster
		if err := dbmarshal.Unmarshal([]byte(dbWorkload.ClusterJSON.String), &cluster); err != nil {
			return model.WorkloadTest{}, fmt.Errorf("failed to unmarshal cluster: %w", err)
		}
		workload.Cluster = &cluster
	}

	return workload, nil
}

// Transaction-based methods for holding locks across operations

// GetAndLockForUpdateInTx gets a workload test with lock within an existing transaction
func (r *workloadTestRepo) GetAndLockForUpdateInTx(
	ctx context.Context,
	tx repo.Tx,
	ID string,
	opts ...options.TableStatement,
) (model.WorkloadTest, error) {
	mysqlTx, ok := tx.(*mysqlTx)
	if !ok {
		return model.WorkloadTest{}, fmt.Errorf("invalid transaction type")
	}

	tblOptions := options.NewTableOption(opts...)
	waitCmd := "NOWAIT"
	if tblOptions.ShouldWait {
		waitCmd = ""
	}

	// Get main workload test record with lock
	var dbWorkload WorkloadTestDB
	query := fmt.Sprintf("SELECT %s FROM workload_tests WHERE id = ? FOR UPDATE %s", workloadTestColumns, waitCmd)
	err := mysqlTx.tx.GetContext(ctx, &dbWorkload, query, ID)
	if err != nil {
		if err == sql.ErrNoRows {
			return model.WorkloadTest{}, fmt.Errorf("workload %s not found", ID)
		}
		return model.WorkloadTest{}, fmt.Errorf("failed to get and lock workload test: %w", err)
	}

	// Convert database record to model
	workload, err := r.dbToModel(dbWorkload)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to convert db record to model: %w", err)
	}

	// Get tasks within the transaction
	tasks, err := r.getTasksInTx(ctx, mysqlTx.tx, ID)
	if err != nil {
		return model.WorkloadTest{}, fmt.Errorf("failed to get tasks: %w", err)
	}
	workload.Tasks = tasks

	return workload, nil
}

// UpdateInTx updates a workload test within an existing transaction
func (r *workloadTestRepo) UpdateInTx(ctx context.Context, tx repo.Tx, workload model.WorkloadTest) error {
	mysqlTx, ok := tx.(*mysqlTx)
	if !ok {
		return fmt.Errorf("invalid transaction type")
	}

	// Serialize params to JSON
	paramsJSON, err := json.Marshal(workload.Params)
	if err != nil {
		return fmt.Errorf("failed to marshal params: %w", err)
	}

	// Serialize cluster to JSON using custom marshaling if it exists
	var clusterJSON sql.NullString
	if workload.Cluster != nil {
		clusterBytes, err := dbmarshal.Marshal(workload.Cluster)
		if err != nil {
			return fmt.Errorf("failed to marshal cluster: %w", err)
		}
		clusterJSON = sql.NullString{String: string(clusterBytes), Valid: true}
	}

	// Update main workload test record - let database handle updated_at
	// Note: started_by is not updated as it should be immutable after creation
	updateQuery := `
		UPDATE workload_tests SET
			name = ?, workload_id = ?, workload_name = ?, params = ?,
			started_at = ?, finished_at = ?, status = ?, message = ?,
			progress = ?, phase = ?, is_auto = ?, cluster = ?
		WHERE id = ?`

	_, err = mysqlTx.tx.ExecContext(ctx, updateQuery,
		workload.Name, workload.WorkloadID, workload.WorkloadName, string(paramsJSON),
		workload.StartedAt, workload.FinishedAt,
		workload.Status, workload.Message,
		workload.Progress, workload.Phase, workload.IsAuto, clusterJSON,
		workload.ID)
	if err != nil {
		return fmt.Errorf("failed to update workload test: %w", err)
	}

	// Upsert tasks within the transaction
	if err := r.upsertTasksInTx(ctx, mysqlTx.tx, workload.ID, workload.Tasks); err != nil {
		return fmt.Errorf("failed to upsert tasks: %w", err)
	}

	return nil
}

// Count returns the count of workload tests group by workload_id
func (r *workloadTestRepo) Count(ctx context.Context) (map[int64]int, error) {
	query := "SELECT workload_id, COUNT(*) as count FROM workload_tests GROUP BY workload_id"

	rows, err := r.db.QueryContext(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("failed to count workload tests by workload_id: %w", err)
	}
	defer rows.Close()

	result := make(map[int64]int)
	for rows.Next() {
		var workloadID int64
		var count int

		if err := rows.Scan(&workloadID, &count); err != nil {
			return nil, fmt.Errorf("failed to scan workload test count: %w", err)
		}

		result[workloadID] = count
	}

	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("failed to iterate workload test count rows: %w", err)
	}

	return result, nil
}

// mysqlTx wraps sqlx.Tx to implement repo.Tx interface
type mysqlTx struct {
	tx *sqlx.Tx
}

func (t *mysqlTx) Commit() error {
	return t.tx.Commit()
}

func (t *mysqlTx) Rollback() error {
	return t.tx.Rollback()
}

// upsertTasksInTx upserts tasks within an existing transaction
func (r *workloadTestRepo) upsertTasksInTx(ctx context.Context, tx *sqlx.Tx, workloadTestID string, tasks []model.WorkloadTask[any]) error {
	if len(tasks) == 0 {
		return nil
	}

	for _, task := range tasks {
		var resultJSON sql.NullString
		if task.Result != nil {
			resultBytes, err := dbmarshal.Marshal(task.Result)
			if err != nil {
				return fmt.Errorf("failed to marshal task result: %w", err)
			}
			resultJSON = sql.NullString{String: string(resultBytes), Valid: true}
		}

		if task.ID > 0 {
			// Update existing task
			updateQuery := `
				UPDATE workload_test_tasks SET
					task_type = ?, started_at = ?, finished_at = ?, status = ?, message = ?,
					progress = ?, related_service_id = ?, phase = ?, result = ?
				WHERE id = ? AND workload_test_id = ?`

			_, err := tx.ExecContext(ctx, updateQuery,
				task.Typ, task.StartedAt, task.FinishedAt,
				task.Status, task.Message, task.Progress,
				task.RelatedServiceID, task.Phase, resultJSON,
				task.ID, workloadTestID)
			if err != nil {
				return fmt.Errorf("failed to update task %d: %w", task.ID, err)
			}
		} else {
			// Insert new task
			insertQuery := `
				INSERT INTO workload_test_tasks (
					workload_test_id, task_type,
					started_at, finished_at, status, message,
					progress, related_service_id, phase, result
				) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`

			_, err := tx.ExecContext(ctx, insertQuery,
				workloadTestID, task.Typ,
				task.StartedAt, task.FinishedAt,
				task.Status, task.Message,
				task.Progress, task.RelatedServiceID, task.Phase, resultJSON)
			if err != nil {
				return fmt.Errorf("failed to insert new task: %w", err)
			}
		}
	}

	return nil
}