package schedule

import (
	"context"
	"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler/schedulers"
	"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler/schedulers/option"
	"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/svc"
	"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/types"
	"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/constants"

	"github.com/zeromicro/go-zero/core/logx"
)

type ScheduleSubmitLogic struct {
	logx.Logger
	ctx    context.Context
	svcCtx *svc.ServiceContext
}

func NewScheduleSubmitLogic(ctx context.Context, svcCtx *svc.ServiceContext) *ScheduleSubmitLogic {
	return &ScheduleSubmitLogic{
		Logger: logx.WithContext(ctx),
		ctx:    ctx,
		svcCtx: svcCtx,
	}
}

func (l *ScheduleSubmitLogic) ScheduleSubmit(req *types.ScheduleReq) (resp *types.ScheduleResp, err error) {
	resp = &types.ScheduleResp{}
	opt := &option.AiOption{
		AdapterId:             req.AiOption.AdapterId,
		TaskName:              req.AiOption.TaskName,
		ResourceType:          req.AiOption.ResourceType,
		Replica:               req.AiOption.Replica,
		Tops:                  req.AiOption.Tops,
		TaskType:              req.AiOption.TaskType,
		DatasetsName:          req.AiOption.Datasets,
		AlgorithmName:         req.AiOption.Algorithm,
		StrategyName:          req.AiOption.Strategy,
		ClusterToStaticWeight: req.AiOption.StaticWeightMap,
		Params:                req.AiOption.Params,
		Envs:                  req.AiOption.Envs,
		Cmd:                   req.AiOption.Cmd,
	}
	aiSchdl, err := schedulers.NewAiScheduler(l.ctx, "", l.svcCtx.Scheduler, opt)
	if err != nil {
		return nil, err
	}

	results, err := l.svcCtx.Scheduler.AssignAndSchedule(aiSchdl)
	if err != nil {
		return nil, err
	}

	switch opt.GetOptionType() {
	case option.AI:
		rs := (results).([]*schedulers.AiResult)
		var synergystatus int64
		if len(rs) > 1 {
			synergystatus = 1
		}
		strategyCode, err := l.svcCtx.Scheduler.AiStorages.GetStrategyCode(req.AiOption.Strategy)

		id, err := l.svcCtx.Scheduler.AiStorages.SaveTask(req.AiOption.TaskName, strategyCode, synergystatus)
		if err != nil {
			return nil, err
		}

		for _, r := range rs {
			scheResult := &types.ScheduleResult{}
			scheResult.ClusterId = r.ClusterId
			scheResult.TaskId = r.TaskId
			scheResult.Strategy = r.Strategy
			scheResult.Replica = r.Replica
			scheResult.Msg = r.Msg
			err := l.svcCtx.Scheduler.AiStorages.SaveAiTask(id, opt, r.ClusterId, r.TaskId, constants.Saved, r.Msg)
			if err != nil {
				return nil, err
			}
			resp.Results = append(resp.Results, scheResult)
		}

	}

	return resp, nil
}