package cn.iocoder.yudao.module.llm.service.async;

import cn.iocoder.yudao.module.llm.dal.dataobject.finetuningtask.FineTuningTaskDO;
import cn.iocoder.yudao.module.llm.dal.dataobject.modelservice.ModelServiceDO;
import cn.iocoder.yudao.module.llm.dal.mysql.basemodel.BaseModelMapper;
import cn.iocoder.yudao.module.llm.dal.mysql.finetuningtask.FineTuningTaskMapper;
import cn.iocoder.yudao.module.llm.dal.mysql.modelservice.ModelServiceMapper;
import cn.iocoder.yudao.module.llm.service.http.TrainHttpService;
import cn.iocoder.yudao.module.llm.service.http.vo.AIgcModelCreateSaveReq;
import cn.iocoder.yudao.module.llm.service.http.vo.AigcModelCreateRespVO;
import cn.iocoder.yudao.module.llm.service.http.vo.AigcModelDeploySaveReq;
import cn.iocoder.yudao.module.llm.service.http.vo.AigcRespVO;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.security.SecureRandom;
import java.util.HashMap;

@Service
public class AsyncModelServiceService {

    private static final Logger log = LoggerFactory.getLogger(AsyncModelServiceService.class);
    @Resource
    private FineTuningTaskMapper fineTuningTaskMapper;

    @Resource
    private ModelServiceMapper modelServiceMapper;

    @Resource
    private TrainHttpService trainHttpService;
    @Resource
    private BaseModelMapper baseModelMapper;

    @Async
    public void createModelService(ModelServiceDO updateObj) {
        AIgcModelCreateSaveReq aIgcModelDeploySaveReq = new AIgcModelCreateSaveReq();
        // 查询基础模型
        /**
         * {"id":9,"baseModelName":"llama3.1",
         * "providerName":"LocalAI",
         * "modelName":"1736432568887qeH05",
         * "modelType":"text-generation",
         * "maxTokens":2048,"isFineTuning":false,
         * "enabled":false,"remark":"","createdAt":"2025-01-09T14:22:48.696991Z",
         * "updatedAt":"2025-01-09T14:22:48.696991Z",
         * "tenants":[],"deployStatus":"",
         * "operation":["edit","delete"],"jobId":"",
         * "lastOperator":"admin","parameters":5,
         * "replicas":1,"label":"","k8sCluster":"",
         * "inferredType":"","gpu":1,"cpu":1,"memory":1,"serviceName":"1736432568887qeh05-9",
         * "containerNames":null,"deployment":{"vllm":false,"status":"","modelPath":"","replicas":0,"inferredType":"","gpu":0,"quantization":""},
         * "fineTuned":null}
         */
        // 获取微调任务
        FineTuningTaskDO fineTuningTaskDO = fineTuningTaskMapper.selectById(updateObj.getFineTuningTask());
        //TODO 先停止 因为 创建 选择基础模型没有部署按钮
//        aIgcModelDeploySaveReq.setBaseModelName(fineTuningTaskDO.getBaseModel());  // 基础模型
        aIgcModelDeploySaveReq.setGpu(updateObj.getGpuCount()); // 设置GPU
        aIgcModelDeploySaveReq.setCpu(1);  // 设置CPU
        aIgcModelDeploySaveReq.setEnabled(true);
        aIgcModelDeploySaveReq.setModelType("text-generation");
        aIgcModelDeploySaveReq.setProviderName("LocalAI");
        aIgcModelDeploySaveReq.setModelName(System.currentTimeMillis()+generateRandomString(5));
        aIgcModelDeploySaveReq.setMaxTokens(4096);
        aIgcModelDeploySaveReq.setParameters(5);
        AigcModelCreateRespVO aigcModelCreateRespVO = trainHttpService.modelCreate(new HashMap<>(), aIgcModelDeploySaveReq);
        log.info("model create info {}",aigcModelCreateRespVO);
        updateObj.setJobId(fineTuningTaskDO.getJobModelListId());
        updateObj.setBaseModelName(fineTuningTaskDO.getJobModelName());
        modelServiceMapper.updateById(updateObj);


    }

    //模型服务开启
    @Async
    public void startModelService(ModelServiceDO updateObj) {
        try {
            /**
             * cpu:1,
             * gpu:0
             * inferredType:"cpu"
             * label:"cpu-aigc-model"
             * maxGpuMemory:0
             * modelWorker:""
             * quantization:"float16"
             * replicas:1
             */
            /**
             * {
             *   "model": "string",
             *   "inferredType": "gpu",
             *   "cpu": 0,
             *   "gpu": 1,
             *   "replicas": 1,
             *   "label": "cpu-aigc-model",
             *   "quantization": "float16",
             *   "maxGpuMemory": 0,
             *   "modelWorker": ""
             * }
             */
            ModelServiceDO modelServiceDO = modelServiceMapper.selectById(updateObj.getId());
            AigcModelDeploySaveReq aigcModelDeploySaveReq = new AigcModelDeploySaveReq(modelServiceDO.getBaseModelName(),
                    1,
                    updateObj.getGpuCount(),
                    "cpu",
                    "cpu-aigc-model",
                    0,
                    "",
                    "float16", 1);
            AigcRespVO aigcRespVO = trainHttpService.modelDeploy(new HashMap<>(), aigcModelDeploySaveReq);
            updateObj.setStatus(1);
            modelServiceMapper.updateById(updateObj);
        }catch(Exception e){
//            updateObj.setStatus(3);
            modelServiceMapper.updateById(updateObj);
        };
    }

    //调型服务关闭
    @Async
    public void stopModelService(ModelServiceDO updateObj) {
        try {
            ModelServiceDO modelServiceDO = modelServiceMapper.selectById(updateObj.getId());
//            Thread.sleep(30000);
            AigcRespVO aigcRespVO = trainHttpService.modelUndeploy(new HashMap<>(), modelServiceDO.getBaseModelName());
            if (aigcRespVO.isSuccess()){
                updateObj.setStatus(0);
            }else {
//                updateObj.setStatus(3);
            }
            modelServiceMapper.updateById(updateObj);
        }catch(Exception e){
//            updateObj.setStatus(3);
            modelServiceMapper.updateById(updateObj);
        };
    }
    private static final String CHARACTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-:";
    private static final SecureRandom RANDOM = new SecureRandom();

    public static String generateRandomString(int length) {
        StringBuilder stringBuilder = new StringBuilder(length);
        for (int i = 0; i < length; i++) {
            int index = RANDOM.nextInt(CHARACTERS.length());
            stringBuilder.append(CHARACTERS.charAt(index));
        }
        return stringBuilder.toString();
    }
}
