package com.yice.webadmin.app.controller;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.github.pagehelper.page.PageMethod;
import com.github.xiaoymin.knife4j.annotations.ApiOperationSupport;
import com.yice.common.core.annotation.MyRequestBody;
import com.yice.common.core.constant.ErrorCodeEnum;
import com.yice.common.core.object.*;
import com.yice.common.core.util.MyCommonUtil;
import com.yice.common.core.util.MyModelUtil;
import com.yice.common.core.util.MyPageUtil;
import com.yice.common.log.annotation.OperationLog;
import com.yice.common.log.model.constant.SysOperationLogType;
import com.yice.webadmin.app.config.LlmModelConfig;
import com.yice.webadmin.app.config.OtherConfig;
import com.yice.webadmin.app.config.PythonConfig;
import com.yice.webadmin.app.dto.ModelDeployDto;
import com.yice.webadmin.app.model.ModelDeploy;
import com.yice.webadmin.app.model.ModelManage;
import com.yice.webadmin.app.model.ModelVersion;
import com.yice.webadmin.app.service.ModelDeployService;
import com.yice.webadmin.app.service.ModelManageService;
import com.yice.webadmin.app.service.ModelVersionService;
import com.yice.webadmin.app.service.ProxyPythonService;
import com.yice.webadmin.app.vo.ModelDeployVo;
import io.swagger.annotations.Api;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

import javax.jws.WebParam;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 模型部署操作控制器类。
 *
 * @author linking
 * @date 2023-04-13
 */
@Api(tags = "模型部署管理接口")
@Slf4j
@RestController
@RequestMapping("/admin/app/modelDeploy")
public class ModelDeployController {

    @Autowired
    private ModelDeployService modelDeployService;
    @Autowired
    private ProxyPythonService proxyPythonService;
    @Autowired
    private OtherConfig otherConfig;
    @Autowired
    private LlmModelConfig llmModelConfig;
    @Autowired
    private PythonConfig pythonConfig;
    @Autowired
    private ModelVersionService modelVersionService;
    @Autowired
    private ModelManageService modelManageService;

    /**
     * 新增模型部署数据。
     *
     * @param modelDeployDto 新增对象。
     * @return 应答结果对象，包含新增对象主键Id。
     */
    @ApiOperationSupport(ignoreParameters = {"modelDeployDto.deployId"})
    @OperationLog(type = SysOperationLogType.ADD)
    @PostMapping("/add")
    public ResponseResult<Long> add(@MyRequestBody ModelDeployDto modelDeployDto) {
        String errorMessage = MyCommonUtil.getModelValidationError(modelDeployDto, false);
        if (errorMessage != null) {
            return ResponseResult.error(ErrorCodeEnum.DATA_VALIDATED_FAILED, errorMessage);
        }
        ModelDeploy modelDeploy = MyModelUtil.copyTo(modelDeployDto, ModelDeploy.class);
        // 验证关联Id的数据合法性
        CallResult callResult = modelDeployService.verifyRelatedData(modelDeploy, null);
        if (!callResult.isSuccess()) {
            return ResponseResult.errorFrom(callResult);
        }
        List<ModelDeploy> modelDeployList = this.modelDeployService.getModelDeployList(modelDeploy, "");
        for (ModelDeploy modelDeploy1 : modelDeployList) {
            if (modelDeploy1.getVersionId() == modelDeploy.getVersionId()) {
                return ResponseResult.error(ErrorCodeEnum.DUPLICATED_UNIQUE_KEY, "同一个模型只能部署一次！");
            }
        }
        ModelVersion modelVersion = this.modelVersionService.getById(modelDeploy.getVersionId());
        ModelManage modelManage = modelManageService.getById(modelVersion.getModelId());
        modelDeploy.setDeployStatus(3);
        modelDeploy.setModelId(modelVersion.getModelId());
        modelDeploy.setVersionName(modelVersion.getVersionName());
        modelDeploy.setModelVersion(modelVersion.getModelVersion());
        modelDeploy.setModelName(modelManage.getModelName());
        modelDeployService.saveNew(modelDeploy);
        return ResponseResult.success(modelDeploy.getDeployId());
    }

    /**
     * 部署指定LLM模型。
     *
     * @return 应答结果对象，包含查询结果集。
     */
    @PostMapping("/deploy")
    public ResponseResult<String> deploy(@MyRequestBody ModelDeployDto modelDeployDto, @MyRequestBody String type) {
        ModelDeploy modelDeploy = modelDeployService.getById(modelDeployDto.getDeployId());
        modelDeploy.setDeployStatus(0);
        modelDeployService.updateById(modelDeploy);
        ResponseResult<String> responseResult = this.doReloadOrStart(modelDeploy, type);
        if (responseResult.isSuccess()) {
            modelDeploy.setDeployStatus(1);
            modelDeployService.updateById(modelDeploy);
            return ResponseResult.success(responseResult.getData());
        }
        modelDeploy.setDeployStatus(-1);
        modelDeployService.updateById(modelDeploy);
        return responseResult;
    }

    private ResponseResult<String> doReloadOrStart(ModelDeploy modelDeploy, String type) {
        String gps_ids = JSON.parseArray(modelDeploy.getResourceInfo()).stream()
                .map(obj -> ((JSONObject) obj).getString("gpu_id"))
                .collect(Collectors.joining(","));
        ModelVersion modelVersion = this.modelVersionService.getById(modelDeploy.getVersionId());
        String requestBody = "{\n" +
                "  \"new_model_name\": \"" + modelDeploy.getVersionName() + "\",\n" +
                "  \"new_model_path\": \"" + modelVersion.getModelUrl() + "\",\n" +
                "  \"startup_param\": {\n" +
                "    \"gpu_ids\": \"" + gps_ids + "\"\n" +
                "  },\n" +
                "  \"controller_address\": \"" + pythonConfig.getControllerAddress() + "\"\n" +
                "}";
        try {
            String apiName = type.equals("start") ? llmModelConfig.getStart() : llmModelConfig.getReload();
            String result = proxyPythonService.predictPost(llmModelConfig.getLlmModelInterface() + apiName, requestBody);
            JSONObject jo = JSON.parseObject(result);
            Integer code = jo.getIntValue("code");
            String msg = jo.getString("msg");
            String data = jo.getString("data");
            if (code != null && code == 200) {
                return ResponseResult.success(data);
            } else {
                return ResponseResult.create(ErrorCodeEnum.SERVER_INTERNAL_ERROR, msg, data);
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 更新模型部署数据。
     *
     * @param modelDeployDto 更新对象。
     * @return 应答结果对象。
     */
    @OperationLog(type = SysOperationLogType.UPDATE)
    @PostMapping("/update")
    public ResponseResult<Void> update(@MyRequestBody ModelDeployDto modelDeployDto) {
        String errorMessage = MyCommonUtil.getModelValidationError(modelDeployDto, true);
        if (errorMessage != null) {
            return ResponseResult.error(ErrorCodeEnum.DATA_VALIDATED_FAILED, errorMessage);
        }
        ModelDeploy modelDeploy = MyModelUtil.copyTo(modelDeployDto, ModelDeploy.class);
        ModelDeploy originalModelDeploy = modelDeployService.getById(modelDeploy.getDeployId());
        if (originalModelDeploy == null) {
            // NOTE: 修改下面方括号中的话述
            errorMessage = "数据验证失败，当前 [数据] 并不存在，请刷新后重试！";
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST, errorMessage);
        }
        // 验证关联Id的数据合法性
        CallResult callResult = modelDeployService.verifyRelatedData(modelDeploy, originalModelDeploy);
        if (!callResult.isSuccess()) {
            return ResponseResult.errorFrom(callResult);
        }
        if (!modelDeployService.update(modelDeploy, originalModelDeploy)) {
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST);
        }
        return ResponseResult.success();
    }

    /**
     * 删除模型部署数据。
     *
     * @param deployId 删除对象主键Id。
     * @return 应答结果对象。
     */
    @OperationLog(type = SysOperationLogType.DELETE)
    @PostMapping("/delete")
    public ResponseResult<Void> delete(@MyRequestBody Long deployId) {
        if (MyCommonUtil.existBlankArgument(deployId)) {
            return ResponseResult.error(ErrorCodeEnum.ARGUMENT_NULL_EXIST);
        }
        return this.doDelete(deployId);
    }

    /**
     * 列出所有可用模型部署列表。
     *
     * @return 应答结果对象，包含查询结果集。
     */
    @PostMapping("/canUseList")
    public ResponseResult<MyPageData<ModelDeployVo>> canUseList() {
        ModelDeploy modelDeployFilter = new ModelDeploy();
        modelDeployFilter.setDeployStatus(1);
        List<ModelDeploy> modelDeployList =
                modelDeployService.getModelDeployListWithRelation(modelDeployFilter, "deploy_id");
        ResponseResult<String> responseResult = this.listRunningModels();
        String data = responseResult.getData();
        String[] modelNames = stringToArray(data);
        List<ModelDeploy> modelDeploys = new ArrayList<ModelDeploy>();
        for (ModelDeploy modelDeploy : modelDeployList) {
            boolean isEx = false;
            for (int i = 0; i < modelNames.length; i++) {
                if (modelDeploy.getVersionName().equals(modelNames[i])) {
                    isEx = true;
                }
            }
            if (isEx) {
                modelDeploy.setDeployStatus(1);
                modelDeploys.add(modelDeploy);
            }
        }
        return ResponseResult.success(MyPageUtil.makeResponseData(modelDeploys, ModelDeploy.INSTANCE));
    }
    /**
     * 列出符合过滤条件的模型部署列表。
     *
     * @param modelDeployDtoFilter 过滤对象。
     * @param orderParam           排序参数。
     * @param pageParam            分页参数。
     * @return 应答结果对象，包含查询结果集。
     */
    @PostMapping("/list")
    public ResponseResult<MyPageData<ModelDeployVo>> list(
            @MyRequestBody ModelDeployDto modelDeployDtoFilter,
            @MyRequestBody MyOrderParam orderParam,
            @MyRequestBody MyPageParam pageParam) {
        if (pageParam != null) {
            PageMethod.startPage(pageParam.getPageNum(), pageParam.getPageSize());
        }
        ModelDeploy modelDeployFilter = MyModelUtil.copyTo(modelDeployDtoFilter, ModelDeploy.class);
        String orderBy = MyOrderParam.buildOrderBy(orderParam, ModelDeploy.class);
        List<ModelDeploy> modelDeployList =
                modelDeployService.getModelDeployListWithRelation(modelDeployFilter, orderBy);
        ResponseResult<String> responseResult = this.listRunningModels();
        if (!responseResult.isSuccess()) {
            modelDeployList.forEach(modelDeploy -> {
                modelDeploy.setDeployStatus(2);
            });
            return ResponseResult.success(MyPageUtil.makeResponseData(modelDeployList, ModelDeploy.INSTANCE));
        }
        String data = responseResult.getData();
        System.out.println(data);
        String[] modelNames = stringToArray(data);
        for (ModelDeploy modelDeploy : modelDeployList) {
            boolean isEx = false;
            for (int i = 0; i < modelNames.length; i++) {
                if (modelDeploy.getVersionName().equals(modelNames[i])) {
                    isEx = true;
                }
            }
            if (!isEx) {
                if (modelDeploy.getDeployStatus() != 3)
                    modelDeploy.setDeployStatus(2);
            } else {
                modelDeploy.setDeployStatus(1);
            }
        }
        return ResponseResult.success(MyPageUtil.makeResponseData(modelDeployList, ModelDeploy.INSTANCE));
    }

    private String[] stringToArray(String strArray) {
        strArray = strArray.replaceAll("\\[|\\]|\"", "");
        return strArray.split(",\\s*"); // 正则表达式也考虑了逗号后的空格
    }

    private ResponseResult<String> listRunningModels() {
        String requestBody = "{\n" +
                "  \"controller_address\": \"" + pythonConfig.getControllerAddress() + "\",\n" +
                "  \"placeholder\": \"string\"\n" +
                "}";
        try {
            String result = proxyPythonService.predictPost(llmModelConfig.getLlmModelInterface() + llmModelConfig.getListRunningModels(), requestBody);
            JSONObject jo = JSON.parseObject(result);
            Integer code = jo.getIntValue("code");
            String msg = jo.getString("msg");
            String data = jo.getString("data");
            if (code != null && code == 200) {
                return ResponseResult.success(data);
            } else {
                return ResponseResult.create(ErrorCodeEnum.SERVER_INTERNAL_ERROR, msg, data);
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 获取GPU信息。
     *
     * @return 应答结果对象，包含查询结果集。
     */
    @PostMapping("/getGpuInfo")
    public ResponseResult<String> getGpuInfo() {
        try {
            String result = proxyPythonService.predictPost(otherConfig.getOtherInterface() + otherConfig.getGetGpuInfo(), "");
            JSONObject jo = JSON.parseObject(result);
            Integer code = jo.getIntValue("code");
            String msg = jo.getString("msg");
            String data = jo.getString("data");
            if (code != null && code == 200) {
                return ResponseResult.success(data);
            } else {
                return ResponseResult.create(ErrorCodeEnum.SERVER_INTERNAL_ERROR, msg, data);
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 停止指定LLM模型。
     *
     * @return 应答结果对象，包含查询结果集。
     */
    @PostMapping("/stop")
    public ResponseResult<String> stop(@MyRequestBody ModelDeployDto modelDeployDto) {
        ModelDeploy modelDeploy = modelDeployService.getById(modelDeployDto.getDeployId());
        ResponseResult<String> responseResult = this.doStop(modelDeploy);
        if (responseResult.isSuccess()) {
            modelDeploy.setDeployStatus(3);
            modelDeployService.updateById(modelDeploy);
            return ResponseResult.success(responseResult.getData());
        }
        return responseResult;
    }

    private ResponseResult<String> doStop(ModelDeploy modelDeploy) {
        String errorMessage = MyCommonUtil.getModelValidationError(modelDeploy, true);
        if (errorMessage != null) {
            return ResponseResult.error(ErrorCodeEnum.DATA_VALIDATED_FAILED, errorMessage);
        }
        String requestBody = "{\n" +
                "  \"model_name\": \"" + modelDeploy.getVersionName() + "\",\n" +
                "  \"controller_address\": \"" + pythonConfig.getControllerAddress() + "\"\n" +
                "}";
        try {
            String result = proxyPythonService.predictPost(llmModelConfig.getLlmModelInterface() + llmModelConfig.getStop(), requestBody);
            JSONObject jo = JSON.parseObject(result);
            Integer code = jo.getIntValue("code");
            String msg = jo.getString("msg");
            String data = jo.getString("data");
            if (code != null && code == 200) {
                return ResponseResult.success(data);
            } else {
                return ResponseResult.create(ErrorCodeEnum.SERVER_INTERNAL_ERROR, msg, data);
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 查看指定模型部署对象详情。
     *
     * @param deployId 指定对象主键Id。
     * @return 应答结果对象，包含对象详情。
     */
    @GetMapping("/view")
    public ResponseResult<ModelDeployVo> view(@RequestParam Long deployId) {
        ModelDeploy modelDeploy = modelDeployService.getByIdWithRelation(deployId, MyRelationParam.full());
        if (modelDeploy == null) {
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST);
        }
        ModelDeployVo modelDeployVo = ModelDeploy.INSTANCE.fromModel(modelDeploy);
        return ResponseResult.success(modelDeployVo);
    }

    private ResponseResult<Void> doDelete(Long deployId) {
        String errorMessage;
        // 验证关联Id的数据合法性
        ModelDeploy originalModelDeploy = modelDeployService.getById(deployId);
        if (originalModelDeploy == null) {
            // NOTE: 修改下面方括号中的话述
            errorMessage = "数据验证失败，当前 [对象] 并不存在，请刷新后重试！";
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST, errorMessage);
        }
        if (!modelDeployService.remove(deployId)) {
            errorMessage = "数据操作失败，删除的对象不存在，请刷新后重试！";
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST, errorMessage);
        }
        this.doStop(originalModelDeploy);
        return ResponseResult.success();
    }
}
