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/imageInference"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/strategy"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/svc"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/types"
	"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/models"
	"net/http"
	"strconv"
)

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

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

//
//func (l *ImageInferenceLogic) ImageInference(req *types.ImageInferenceReq) (resp *types.ImageInferenceResp, err error) {
//	return nil, nil
//}

func (l *ImageInferenceLogic) ImageInfer(r *http.Request, req *types.ImageInferenceReq) (resp *types.ImageInferenceResp, err error) {
	resp = &types.ImageInferenceResp{}
	if len(req.InstanceIds) == 0 {
		return nil, errors.New("instances are empty")
	}

	var instanceList []*models.AiInferDeployInstance
	for _, id := range req.InstanceIds {
		instance, err := l.svcCtx.Scheduler.AiStorages.GetInferDeployInstanceById(id)
		if err != nil {
			return nil, err
		}
		if instance == nil {
			return nil, errors.New("instance is empty ")
		}
		instanceList = append(instanceList, instance)
	}

	if len(instanceList) == 0 {
		return nil, errors.New("instanceList are empty")
	}

	// process uploaded images
	var ts []*imageInference.ImageFile

	uploadedFiles := r.MultipartForm.File

	if len(uploadedFiles) == 0 {
		return nil, errors.New("Images does not exist")
	}

	if len(uploadedFiles["images"]) == 0 {
		return nil, errors.New("Images does not exist")
	}

	for _, header := range uploadedFiles["images"] {
		file, err := header.Open()
		if err != nil {
			return nil, err
		}
		defer file.Close()
		var ir types.ImageResult
		ir.ImageName = header.Filename
		t := imageInference.ImageFile{
			ImageResult: &ir,
			File:        file,
		}
		ts = append(ts, &t)
	}

	//single adapter logic
	if len(req.StaticWeightMap) != 1 {
		return nil, errors.New("staticWeightMap != 1")
	}

	adapterId := strconv.FormatInt(instanceList[0].AdapterId, 10)
	staticWeightMap, ok := req.StaticWeightMap[adapterId]
	if !ok {
		return nil, errors.New("set staticWeightMap failed")
	}

	// create InferOption
	opt := &option.InferOption{
		TaskName:  req.TaskName,
		TaskDesc:  req.TaskDesc,
		AdapterId: adapterId,
		//AiClusterIds:    req.AiClusterIds,
		//ModelName:       req.ModelName,
		ModelType:       req.ModelType,
		Strategy:        req.Strategy,
		StaticWeightMap: staticWeightMap,
	}

	adapterName, err := l.svcCtx.Scheduler.AiStorages.GetAdapterNameById(opt.AdapterId)
	if err != nil {
		return nil, err
	}

	// set strategy
	if opt.Strategy == "" {
		return nil, errors.New("strategy is empty")
	}

	var strat strategy.Strategy
	switch opt.Strategy {
	case strategy.STATIC_WEIGHT:
		strat = strategy.NewStaticWeightStrategy(opt.StaticWeightMap, int32(len(ts)))
		if err != nil {
			return nil, err
		}
	default:
		return nil, errors.New("no strategy has been chosen")
	}
	clusters, err := strat.Schedule()
	if err != nil {
		return nil, err
	}

	if clusters == nil || len(clusters) == 0 {
		return nil, errors.New("clusters is nil")
	}

	//remove empty replica
	for i := len(clusters) - 1; i >= 0; i-- {
		if clusters[i].Replicas == 0 {
			clusters = append(clusters[:i], clusters[i+1:]...)
		}
	}

	// create inference struct
	imageInfer, err := imageInference.New(imageInference.NewImageClassification(), ts, clusters, instanceList, opt, l.svcCtx.Scheduler.AiStorages, l.svcCtx.Scheduler.AiService.InferenceAdapterMap, adapterName)
	if err != nil {
		return nil, err
	}

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

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

	go func() {
		err := in.In.InferTask(id)
		if err != nil {
			logx.Errorf(err.Error())
			return
		}
	}()

	return resp, nil
}