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"

	"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,
		ResourceType:          req.AiOption.ResourceType,
		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)
		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
			resp.Results = append(resp.Results, scheResult)
		}
		err = l.svcCtx.Scheduler.AiStorages.SaveTask(req.AiOption.TaskName)
		if err != nil {
			return nil, err
		}
	}

	return resp, nil
}