package gateway

import (
	"context"
	"encoding/json"
	"fmt"
	"strconv"

	"github.com/tidbcloud/pov-service/internal/client/rms"
	"github.com/tidbcloud/pov-service/internal/client/tcms"
	"github.com/tidbcloud/pov-service/internal/client/tiworkload"
	"github.com/tidbcloud/pov-service/internal/common/params"
	"github.com/tidbcloud/pov-service/internal/config"
	"github.com/tidbcloud/pov-service/internal/model"
	"github.com/tidbcloud/pov-service/internal/pkg/log"
)

const (
	// we use different workflow for creating cluster in test and prod environment
	tcmsCreateNGClusterWorkflowID     = 120036
	tcmsCreateNGProdClusterWorkflowID = 120046

	tcmsCreateEssentialClusterWorkflowID = 150090

	// defaultDBUser = "root"
	// defaultDBPwd  = "12345678"

	defOverProvisionTempl    = "pov-%s"
	defOverProvisionReplicas = 5

	// tcmsCreateClusterWorkflowID = 120041 // workflow for creating ksyun testbed, used in test phase
)

var (
	defOverProvisionResource = tiworkload.OverProvisionResource{
		CPU:    "7",
		Memory: "13Gi",
	}
)

// WorkloadGateway defines the interface for workload operations
type Workload interface {
	ListWorkloads(params.ListWorkloads) ([]model.Workload, error)
	TriggerCreateCluster(ctx context.Context, params model.WorkloadParams) (model.WorkloadTaskBase, error)
	PollCreateClusterTask(
		ctx context.Context,
		tidbCluster *model.TiDBCluster,
		task model.WorkloadTaskBase,
	) (model.WorkloadTaskBase, model.CreateClusterResult, error)
	TriggerPrepareData(
		ctx context.Context,
		tidbCluster model.TiDBCluster,
		params params.RunWorkload) (model.WorkloadTaskBase, error)
	PollPrepareDataTask(
		ctx context.Context,
		tidbCluster model.TiDBCluster,
		task model.WorkloadTaskBase,
	) (model.WorkloadTaskBase, model.PrepareDataResult, error)
	TriggerRunTraffic(
		ctx context.Context,
		tidbCluster model.TiDBCluster,
		prepareResult model.PrepareDataResult,
		params model.WorkloadParams) (model.WorkloadTaskBase, error)
	PollRunTrafficTask(
		ctx context.Context,
		task model.WorkloadTaskBase,
	) (model.WorkloadTaskBase, error)
	ReleaseCluster(
		ctx context.Context,
		cluster model.TiDBCluster,
	) error
	CancelTask(
		ctx context.Context,
		task model.WorkloadTaskBase,
	) error
	ScaleCluster(
		ctx context.Context,
		cluster model.TiDBCluster,
		replicas int,
	) error
	UpdateClusterRU(
		ctx context.Context,
		cluster model.TiDBCluster,
		minRU, maxRU int64,
	) error
	InjectFailure(
		ctx context.Context,
		cluster model.TiDBCluster,
		duration string,
	) error
	RollbackFailure(
		ctx context.Context,
		cluster model.TiDBCluster,
	) error
	TriggerRecreateTraffic(
		ctx context.Context,
		oldTrafficID int64,
		newparams model.WorkloadParams) (model.WorkloadTaskBase, error)
	SampleDBs(
		ctx context.Context,
		cluster model.TiDBCluster,
		options model.SampleDBOptions) (model.DBSampleResult, error)
	SampleData(
		ctx context.Context,
		cluster model.TiDBCluster,
		options model.SampleDataOptions) (model.TableQueryResult, error)
	OverProvision(
		ctx context.Context,
		cluster model.TiDBCluster,
	) error
	DeleteOverProvision(
		ctx context.Context,
		cluster model.TiDBCluster,
	) error
	ExecuteSQL(
		ctx context.Context,
		cluster model.TiDBCluster,
		sql string,
	) error
}

type workloadImpl struct {
	tcmsClient       tcms.Client
	rmsClient        rms.Client
	tiworkloadClient tiworkload.Client
	paramHelper      paramHelper
}

// Ensure Workload implements WorkloadGateway interface
var _ Workload = (*workloadImpl)(nil)

func NewWorkload(
	tcmsClient tcms.Client,
	rmsClient rms.Client,
	tiworkloadClient tiworkload.Client) Workload {
	return workloadImpl{
		tcmsClient:       tcmsClient,
		rmsClient:        rmsClient,
		tiworkloadClient: tiworkloadClient,
		paramHelper: newParamHelper(
			tiworkloadClient,
			config.IsProdEnv()),
	}
}

func (w workloadImpl) ListWorkloads(param params.ListWorkloads) ([]model.Workload, error) {
	povWorkflows, err := w.tcmsClient.ListPovWorkflows()
	if err != nil {
		return nil, err
	}
	all := make([]model.Workload, 0, len(povWorkflows))
	for _, workflow := range povWorkflows {
		var (
			workload model.Workload
		)
		// INFO: for demo phase, we use workflow description to store workload metadata
		// if it's json unmarshalable, we just use the json info stored in workflow
		// otherwise fallback to use aggregated workflow info from workflow/tiworkload
		if err := json.Unmarshal([]byte(workflow.Desc), &workload); err != nil {
			workload = model.Workload{
				Description: workflow.Desc,
			}

			tiwls := workflow.TrafficNames()
			tags := make([]string, 0)
			for _, wl := range tiwls {
				wl, err := w.tiworkloadClient.OneTrafficByName(wl)
				if err != nil {
					return nil, err
				}
				tags = append(tags, wl.Feature.Tags...)
				workload.Type = wl.WorkloadType()
			}

			workload.Name = workflow.Name
			workload.Tags = tags
		}

		workload.ID = workflow.ID
		workload.ParamsTemplate = workflow.Spec.StartNodeParams

		all = append(all, workload)
	}

	return all, nil
}

func (w workloadImpl) TriggerCreateCluster(
	ctx context.Context,
	params model.WorkloadParams,
) (model.WorkloadTaskBase, error) {
	workloadID, err := w.paramHelper.clusterWorkflowID(params)
	if err != nil {
		return model.WorkloadTaskBase{}, err
	}

	// TODO: set non_block_conn to false if the cluster should be prepared using sql generate
	workflowParams := tcms.NewRunWorkflowParams(
		int64(workloadID),
		"tidbcloud_pov_cluster_create",
		params,
	)
	if runnerID, err := w.paramHelper.tcmsRunnerID(params); err == nil {
		workflowParams.RunnerID = runnerID
	} else {
		return model.WorkloadTaskBase{}, err
	}

	workflow, err := w.tcmsClient.RunWorkflow(ctx, *workflowParams)
	if err != nil {
		return model.WorkloadTaskBase{}, err
	}

	task := *workflow.Convert2WorkloadTask()
	task.Typ = model.WorkloadTaskTypeCreateCluster

	return task, nil
}

// PollCreateClusterTask polls the status and result of a create cluster task
func (w workloadImpl) PollCreateClusterTask(
	ctx context.Context,
	tidbCluster *model.TiDBCluster, // create cluster may be a multi-step process, so we'll need the intermediate result
	task model.WorkloadTaskBase,
) (model.WorkloadTaskBase, model.CreateClusterResult, error) {
	var result model.CreateClusterResult
	if tidbCluster == nil {
		tidbCluster = &model.TiDBCluster{}
	}
	result.Cluster = tidbCluster

	if tidbCluster == nil || tidbCluster.IsCreating() {
		log.Log.Infof("Polling create cluster from workflow for task %s", task.ID)
		id, err := strconv.ParseInt(task.RelatedServiceID, 10, 64)
		if err != nil {
			return model.WorkloadTaskBase{}, model.CreateClusterResult{}, fmt.Errorf("invalid related service ID: %w", err)
		}
		workflowRun, err := w.tcmsClient.GetWorkflowRun(ctx, id)
		if err != nil {
			return model.WorkloadTaskBase{}, model.CreateClusterResult{}, err
		}
		workflowTask := workflowRun.Convert2WorkloadTask()

		workflowTask.Typ = model.WorkloadTaskTypeCreateCluster

		isComplete := workflowTask.Status == model.ExecStatusCompleted

		// If task is completed, extract the cluster result
		// TODO: verify the outputs/StateOutput of the task
		if isComplete {
			createTask := workflowRun.GetFinalizedTask("TESTBED_CREATE")
			cloudCluster, err := tcms.NewCreateCloudClusterOutput(createTask)
			if err != nil {
				return model.WorkloadTaskBase{}, model.CreateClusterResult{}, err
			}

			if result.Cluster == nil {
				result.Cluster = &model.TiDBCluster{}
			}
			result.Cluster.TestbedName = cloudCluster.Name
			result.Cluster.TestbedID = cloudCluster.ID
			result.Cluster.TestbedPlanExecID = cloudCluster.PlanExecID
			result.Cluster.Host = cloudCluster.Host
			result.Cluster.Port = cloudCluster.Port
			result.Cluster.TiDBCloudClusterID = cloudCluster.TiDBCloudClusterID
			result.Cluster.TiDBCloudOrgID = cloudCluster.TiDBCloudOrg
			result.Cluster.TiDBCloudProjectID = cloudCluster.TiDBCloudProject
			result.Cluster.MaxRU = cloudCluster.MaxRU
			result.Cluster.MinRU = cloudCluster.MinRU
			result.Cluster.Type = cloudCluster.ClusterType
		} else {
			task.Status = workflowTask.Status
			return task, result, nil
		}
	}

	if result.Cluster.ShouldRegister2TiWorkload() {
		log.Log.Infof("Registering testbed %s to TiWorkload", result.Cluster.TestbedName)
		if agentID, err := w.tiworkloadClient.RegisterTestbed(ctx, result.Cluster.TestbedName); err == nil {
			result.Cluster.TiWorkloadAgentID = agentID
			return task, result, nil
		} else {
			return model.WorkloadTaskBase{}, model.CreateClusterResult{}, err
		}
	}

	if result.Cluster.IsRegisting2TiWorkload() {
		if clstID, err := w.tiworkloadClient.Get1stClusterByAgent(ctx, result.Cluster.TiWorkloadAgentID); err == nil && clstID != "" {
			result.Cluster.TiWorkloadClusterID = clstID
			task.Status = model.ExecStatusCompleted
			return task, result, nil
		}
	}

	return task, result, nil
}

func (w workloadImpl) TriggerPrepareData(
	ctx context.Context,
	tidbCluster model.TiDBCluster,
	params params.RunWorkload,
) (model.WorkloadTaskBase, error) {
	workloadName := params.Params.GetWorkloadName()
	if workloadName == "" {
		return model.WorkloadTaskBase{}, fmt.Errorf("workload name is empty")
	}

	workloadParams, err := w.paramHelper.prepareParams(workloadName, params.GetParams())
	if err != nil {
		return model.WorkloadTaskBase{}, err
	}

	dbHost := tidbCluster.Host
	dbPort := tidbCluster.Port

	tiworkloadParams := tiworkload.PrepareParams{
		ExecuteBase: tiworkload.ExecuteBase{
			ClusterID: tidbCluster.TiWorkloadClusterID,
			DBHost:    dbHost,
			DBPort:    dbPort,
			// configable if needed

			WorkloadName: workloadName,
		},
		PrepareParams: tiworkload.TaskParams{Params: workloadParams},
	}

	id, err := w.tiworkloadClient.RunPrepare(tiworkloadParams)
	if err != nil {
		return model.WorkloadTaskBase{}, err
	}

	return model.WorkloadTaskBase{
		Typ:              model.WorkloadTaskTypePrepareData,
		ExecBase:         model.ExecBase{Status: model.ExecStatusRunning},
		Progress:         0,
		RelatedServiceID: strconv.FormatInt(id, 10),
	}, nil
}

// PollPrepareDataTask polls the status and result of a prepare data task
func (w workloadImpl) PollPrepareDataTask(
	ctx context.Context,
	tidbCluster model.TiDBCluster,
	task model.WorkloadTaskBase,
) (ret model.WorkloadTaskBase, result model.PrepareDataResult, err error) {
	// check whether the preparation has finished
	id, err := strconv.ParseInt(task.RelatedServiceID, 10, 64)
	if err != nil {
		return
	}

	exec, err := w.tiworkloadClient.GetTrafficExec(id)
	if err != nil {
		return
	}

	// If task is not under stopping phase and status is stopped,
	// the task is terminated by downstream svces, which the user couldn't know.
	// In this case, the task should be set to failed
	ret = exec.ToWorkloadTask()
	ret.Typ = model.WorkloadTaskTypePrepareData

	if ret.IsAbnormalFinished() {
		return ret, result, nil
	}

	// check whether the privatelink has been created
	if tidbCluster.HasConnection() {
		return
	} else {
		clst, err := w.tiworkloadClient.GetCluster(ctx, tidbCluster.TiWorkloadClusterID)
		if err != nil {
			return model.WorkloadTaskBase{}, model.PrepareDataResult{}, err
		}

		if clst.DBHost != "" && clst.DBPort != 0 {
			ret.Progress = (ret.Progress + 1) / 2
			result.DBHost = clst.DBHost
			result.DBPort = clst.DBPort
		} else {
			ret.Progress = (ret.Progress + 0) / 2
			ret.Status = model.ExecStatusRunning
		}
	}

	return ret, result, nil
}

func (w workloadImpl) TriggerRunTraffic(
	ctx context.Context,
	tidbCluster model.TiDBCluster,
	_ model.PrepareDataResult,
	params model.WorkloadParams,
) (model.WorkloadTaskBase, error) {
	workloadName := params.GetWorkloadName()

	workloadParams, err := w.paramHelper.trafficParams(workloadName, params)
	if err != nil {
		return model.WorkloadTaskBase{}, err
	}

	dbHost := tidbCluster.Host
	dbPort := tidbCluster.Port

	tiworkloadParams := tiworkload.WorkloadParams{
		ExecuteBase: tiworkload.ExecuteBase{
			ClusterID: tidbCluster.TiWorkloadClusterID,
			DBHost:    dbHost,
			DBPort:    dbPort,

			WorkloadName: workloadName,
		},
		WorkloadParams: tiworkload.TaskParams{Params: workloadParams},
	}

	id, err := w.tiworkloadClient.RunTraffic(tiworkloadParams)
	if err != nil {
		return model.WorkloadTaskBase{}, err
	}

	return model.WorkloadTaskBase{
		Typ:              model.WorkloadTaskTypeRunTraffic,
		ExecBase:         model.ExecBase{Status: model.ExecStatusRunning},
		Progress:         0,
		RelatedServiceID: strconv.FormatInt(id, 10),
	}, nil
}

// PollRunTrafficTask polls the status and result of a run traffic task
func (w workloadImpl) PollRunTrafficTask(
	ctx context.Context,
	task model.WorkloadTaskBase,
) (model.WorkloadTaskBase, error) {
	id, err := strconv.ParseInt(task.RelatedServiceID, 10, 64)
	if err != nil {
		return model.WorkloadTaskBase{}, fmt.Errorf("invalid related service ID: %w", err)
	}

	trafficExec, err := w.tiworkloadClient.GetTrafficExec(id)
	if err != nil {
		return model.WorkloadTaskBase{}, err
	}

	// If task is not under stopping phase and status is stopped,
	// the task is terminated by downstream svces, which the user couldn't know.
	// In this case, the task should be set to failed
	updatedTask := trafficExec.ToWorkloadTask()
	updatedTask.Typ = model.WorkloadTaskTypeRunTraffic

	return updatedTask, nil
}

func (w workloadImpl) ReleaseCluster(
	ctx context.Context,
	cluster model.TiDBCluster,
) error {
	// return nil
	return w.tcmsClient.ReleaseTestbed(ctx, cluster.TestbedPlanExecID, cluster.TestbedName)
}

func (w workloadImpl) CancelTask(
	ctx context.Context,
	task model.WorkloadTaskBase,
) error {
	if task.RelatedServiceID == "" {
		return fmt.Errorf("invalid related service ID")
	}
	switch task.Typ {
	case model.WorkloadTaskTypeCreateCluster, model.WorkloadTaskTypePrepareData:
		if v, err := strconv.ParseInt(task.RelatedServiceID, 10, 64); err == nil {
			return w.tcmsClient.CancelWorkflow(ctx, v)
		} else {
			return err
		}
	case model.WorkloadTaskTypeRunTraffic:
		if v, err := strconv.ParseInt(task.RelatedServiceID, 10, 64); err == nil {
			return w.tiworkloadClient.TerminateTrafficExec(v)
		} else {
			return err
		}

	}
	return nil
}

func (w workloadImpl) ScaleCluster(
	ctx context.Context,
	cluster model.TiDBCluster,
	replicas int,
) error {
	return w.rmsClient.ScaleTiDBGroup(ctx, cluster, rms.ScaleTiDBGroupParams{
		Replicas: replicas,
	})
}

func (w workloadImpl) UpdateClusterRU(
	ctx context.Context,
	cluster model.TiDBCluster,
	minRU, maxRU int64,
) error {
	return w.rmsClient.UpdateTiDBGroupRU(ctx, cluster, minRU, maxRU)
}

func (w workloadImpl) InjectFailure(
	ctx context.Context,
	cluster model.TiDBCluster,
	duration string,
) error {
	return w.rmsClient.InjectTiDBGroupFailure(ctx, cluster, rms.InjectTiDBGroupFailureParams{
		FailureFault: &rms.InjectFailureFaultParams{
			Duration: duration,
		},
	})
}

func (w workloadImpl) RollbackFailure(
	ctx context.Context,
	cluster model.TiDBCluster,
) error {
	return w.rmsClient.RollbackTiDBGroupFailure(ctx, cluster, rms.RollbackTiDBGroupFailureParams{})
}

func (w workloadImpl) TriggerRecreateTraffic(
	ctx context.Context,
	oldTrafficID int64,
	newparams model.WorkloadParams) (model.WorkloadTaskBase, error) {
	trafficExec, err := w.tiworkloadClient.GetTrafficExec(oldTrafficID)
	if err != nil {
		return model.WorkloadTaskBase{}, err
	}
	params, err := w.paramHelper.trafficParams(trafficExec.Workload, newparams)
	if err != nil {
		return model.WorkloadTaskBase{}, err
	}

	dbHost := trafficExec.DBHost
	dbPort := trafficExec.DBPort

	tiworkloadParams := tiworkload.WorkloadParams{
		ExecuteBase: tiworkload.ExecuteBase{
			ClusterID: trafficExec.ClusterId,
			DBHost:    dbHost,
			DBPort:    dbPort,
			// configable if needed
			User:         trafficExec.User,
			Password:     trafficExec.Password,
			WorkloadName: trafficExec.Workload,
		},

		WorkloadParams: tiworkload.TaskParams{Params: params},
	}

	newid, err := w.tiworkloadClient.RecreateTraffic(ctx, oldTrafficID, tiworkloadParams)
	if err != nil {
		return model.WorkloadTaskBase{}, err
	}

	return model.WorkloadTaskBase{
		Typ:              model.WorkloadTaskTypeRunTraffic,
		ExecBase:         model.ExecBase{Status: model.ExecStatusRunning},
		Progress:         0,
		RelatedServiceID: strconv.FormatInt(newid, 10),
	}, nil
}

func (w workloadImpl) SampleDBs(
	ctx context.Context,
	cluster model.TiDBCluster,
	options model.SampleDBOptions) (model.DBSampleResult, error) {
	if cluster.TiWorkloadClusterID == "" {
		return model.DBSampleResult{}, fmt.Errorf("cluster is not ready for sampling")
	}
	ret, err := w.tiworkloadClient.SampleDBs(
		ctx, tiworkload.DBAccess{
			ClusterID: cluster.TiWorkloadClusterID,
		}, options,
	)
	if err != nil {
		return model.DBSampleResult{}, err
	}
	return *ret, nil
}

func (w workloadImpl) SampleData(
	ctx context.Context,
	cluster model.TiDBCluster,
	options model.SampleDataOptions) (model.TableQueryResult, error) {
	if cluster.TiWorkloadClusterID == "" {
		return model.TableQueryResult{}, fmt.Errorf("cluster is not ready for sampling")
	}
	ret, err := w.tiworkloadClient.SampleData(
		ctx, tiworkload.DBAccess{
			ClusterID: cluster.TiWorkloadClusterID,
		}, options,
	)
	if err != nil {
		return model.TableQueryResult{}, err
	}
	return *ret, nil
}

func (w workloadImpl) OverProvision(
	ctx context.Context,
	cluster model.TiDBCluster,
) error {
	defaultName := fmt.Sprintf(defOverProvisionTempl, cluster.TiWorkloadClusterID)

	_, err := w.tiworkloadClient.PostOverProvision(ctx, cluster.TiWorkloadClusterID,
		tiworkload.OverProvisionOption{
			Name:     defaultName,
			Resource: defOverProvisionResource,
			Replicas: defOverProvisionReplicas,
		})
	return err
}

func (w workloadImpl) DeleteOverProvision(
	ctx context.Context,
	cluster model.TiDBCluster,
) error {
	defaultName := fmt.Sprintf(defOverProvisionTempl, cluster.TiWorkloadClusterID)

	return w.tiworkloadClient.DeleteOverprovision(ctx, cluster.TiWorkloadClusterID, defaultName)

}

func (w workloadImpl) ExecuteSQL(
	ctx context.Context,
	cluster model.TiDBCluster,
	sql string,
) error {
	_, err := w.tiworkloadClient.ExecuteSql(ctx,
		tiworkload.ExecteSqlReq{
			ClusterID: cluster.TiWorkloadClusterID,
			Query:     sql,
		})
	return err
}