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"
	"net/http"
)

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{}
	opt := &option.InferOption{
		TaskName:        req.TaskName,
		TaskDesc:        req.TaskDesc,
		AdapterId:       req.AdapterId,
		AiClusterIds:    req.AiClusterIds,
		ModelName:       req.ModelName,
		ModelType:       req.ModelType,
		Strategy:        req.Strategy,
		StaticWeightMap: req.StaticWeightMap,
	}

	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)
	}

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

	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")
	}

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

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

	imageInfer, err := imageInference.New(imageInference.NewImageClassification(), ts, clusters, 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
}