package inference

import (
	"context"
	"errors"
	"github.com/zeromicro/go-zero/core/logx"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/schedulers/option"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/service/inference"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/service/inference/textInference"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/svc"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/types"
)

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

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

func (l *TextToTextInferenceLogic) TextToTextInference(req *types.TextToTextInferenceReq) (resp *types.TextToTextInferenceResp, err error) {
	resp = &types.TextToTextInferenceResp{}
	opt := &option.InferOption{
		TaskName:     req.TaskName,
		TaskDesc:     req.TaskDesc,
		AdapterId:    req.AdapterId,
		AiClusterIds: req.AiClusterIds,
		ModelName:    req.ModelName,
		ModelType:    req.ModelType,
	}

	_, ok := l.svcCtx.Scheduler.AiService.AiCollectorAdapterMap[opt.AdapterId]
	if !ok {
		return nil, errors.New("AdapterId does not exist")
	}

	adapterName, err := l.svcCtx.Scheduler.AiStorages.GetAdapterNameById(opt.AdapterId)

	inType, err := textInference.NewTextToText(opt, l.svcCtx.Scheduler.AiStorages, l.svcCtx.Scheduler.AiService.InferenceAdapterMap)
	if err != nil {
		return nil, err
	}
	textInfer, err := textInference.New(inType, opt, l.svcCtx.Scheduler.AiStorages, l.svcCtx.Scheduler.AiService.InferenceAdapterMap, adapterName)
	if err != nil {
		return nil, err
	}

	in := inference.Inference{
		In: textInfer,
	}

	id, err := in.In.CreateTask()
	if err != nil {
		return nil, err
	}
	err = in.In.InferTask(id)
	if err != nil {
		return nil, err
	}

	return resp, nil
}
