package com.yice.webadmin.app.controller;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONException;
import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.databind.ObjectMapper;
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.annotation.NoAuthInterface;
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.PythonConfig;
import com.yice.webadmin.app.dto.RunPublishDto;
import com.yice.webadmin.app.dto.TuningRunDto;
import com.yice.webadmin.app.model.DatasetVersion;
import com.yice.webadmin.app.model.ModelVersion;
import com.yice.webadmin.app.model.TuningRun;
import com.yice.webadmin.app.service.*;
import com.yice.webadmin.app.util.FileUtil;
import com.yice.webadmin.app.vo.TuningRunVo;
import io.swagger.annotations.Api;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.Resource;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;

import java.io.File;
import java.io.IOException;
import java.util.Arrays;
import java.util.Date;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * 精调任务运行操作控制器类。
 *
 * @author linking
 * @date 2023-04-13
 */
@Api(tags = "精调任务运行管理接口")
@Slf4j
@RestController
@RequestMapping("/admin/app/tuningRun")
public class TuningRunController {

    @Autowired
    private TuningRunService tuningRunService;
    @Autowired
    private TuningTaskService tuningTaskService;
    @Autowired
    private ModelVersionService modelVersionService;
    @Autowired
    private ModelTaskService modelTaskService;
    @Autowired
    private DatasetVersionService datasetVersionService;
    @Autowired
    private PythonConfig pythonConfig;

    /**
     * 新增精调任务运行数据。
     *
     * @param tuningRunDto 新增对象。
     * @return 应答结果对象，包含新增对象主键Id。
     */
    @ApiOperationSupport(ignoreParameters = {"tuningRunDto.runId"})
    @OperationLog(type = SysOperationLogType.ADD)
    @PostMapping("/add")
    public ResponseResult<Long> add(@MyRequestBody TuningRunDto tuningRunDto) {
        String errorMessage = MyCommonUtil.getModelValidationError(tuningRunDto, false);
        if (errorMessage != null) {
            return ResponseResult.error(ErrorCodeEnum.DATA_VALIDATED_FAILED, errorMessage);
        }
        if (tuningRunDto.getTaskId() == null) {
            return ResponseResult.error(ErrorCodeEnum.DATA_VALIDATED_FAILED, "请填写任务ID");
        }
        TuningRun tuningRun = MyModelUtil.copyTo(tuningRunDto, TuningRun.class);
        tuningRun = tuningRunService.saveNew(tuningRun);
        return ResponseResult.success(tuningRun.getRunId());
    }

    /**
     * 更新精调任务运行数据。
     *
     * @param tuningRunDto 更新对象。
     * @return 应答结果对象。
     */
    @OperationLog(type = SysOperationLogType.UPDATE)
    @PostMapping("/update")
    public ResponseResult<Void> update(@MyRequestBody TuningRunDto tuningRunDto) {
        String errorMessage = MyCommonUtil.getModelValidationError(tuningRunDto, true);
        if (errorMessage != null) {
            return ResponseResult.error(ErrorCodeEnum.DATA_VALIDATED_FAILED, errorMessage);
        }
        TuningRun tuningRun = MyModelUtil.copyTo(tuningRunDto, TuningRun.class);
        TuningRun originalTuningRun = tuningRunService.getById(tuningRun.getRunId());
        if (originalTuningRun == null) {
            // NOTE: 修改下面方括号中的话述
            errorMessage = "数据验证失败，当前 [数据] 并不存在，请刷新后重试！";
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST, errorMessage);
        }
        if (!tuningRunService.update(tuningRun, originalTuningRun)) {
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST);
        }
        return ResponseResult.success();
    }

    /**
     * 更新精调任务运行数据,通过ID。
     *
     * @param tuningRunDto 更新对象。
     * @return 应答结果对象。
     */
    @OperationLog(type = SysOperationLogType.UPDATE)
    @PostMapping("/updateById")
    public ResponseResult<Void> updateById(@MyRequestBody TuningRunDto tuningRunDto) {
        String errorMessage = MyCommonUtil.getModelValidationError(tuningRunDto, true);
        if (errorMessage != null) {
            return ResponseResult.error(ErrorCodeEnum.DATA_VALIDATED_FAILED, errorMessage);
        }
        TuningRun tuningRun = MyModelUtil.copyTo(tuningRunDto, TuningRun.class);
        TuningRun originalTuningRun = tuningRunService.getById(tuningRun.getRunId());
        if (originalTuningRun == null) {
            errorMessage = "数据验证失败，当前 [数据] 并不存在，请刷新后重试！";
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST, errorMessage);
        }
        originalTuningRun.setRunStatus(tuningRun.getRunStatus());
        if (!tuningRunService.updateById(originalTuningRun)) {
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST);
        }
        return ResponseResult.success();
    }

    /**
     * 获取预览命令。
     *
     * @param runId 运行ID。
     * @return 应答结果对象。
     */
    @GetMapping("/getPreviewCommand")
    public ResponseResult<JSONArray> getPreviewCommand(@RequestParam Long runId) {
        TuningRun tuningRun = this.tuningRunService.getById(runId);
        ModelVersion modelVersion = this.modelVersionService.getById(tuningRun.getModelVersionId());
        JSONObject jsonObject = (JSONObject) JSON.parse(tuningRun.getConfiguration());
        JSONArray datasetVersionNames = new JSONArray();
        Set<Long> datasetVersionIds = Arrays.stream(tuningRun.getDatasetVersionIds().split(",")).mapToLong(Long::parseLong).boxed().collect(Collectors.toSet());
        datasetVersionService.getInList(datasetVersionIds).stream()
                .map(DatasetVersion::getVersionName)
                .forEach(versionName -> {
                    try {
                        datasetVersionNames.add(versionName);
                    } catch (JSONException e) {
                        e.printStackTrace();
                    }
                });
        JSONArray array = new JSONArray();
        array.add("zh");
        array.add(modelVersion.getVersionName());
        array.add(modelVersion.getModelUrl());
        array.add(tuningRun.getTrainMethod());
        array.add(new JSONArray());
        array.add(jsonObject.get("quantizationLevel"));
        array.add(jsonObject.get("promptTemplate"));
        array.add("");
        array.add(false);
        array.add(false);
        array.add("none");
        array.add("Supervised Fine-Tuning");
        array.add(pythonConfig.getDatasetFileMenu());
        array.add(datasetVersionNames);
        array.add(jsonObject.get("truncationLength"));
        array.add(jsonObject.get("learningRate"));
        array.add(jsonObject.get("iterationRound"));
        array.add(jsonObject.get("maximumSampleSize"));
        array.add(jsonObject.get("CalculationType"));
        array.add(jsonObject.get("batchSize"));
        array.add(jsonObject.get("gradientAccumulation"));
        array.add(jsonObject.get("LearningRateRegulator"));
        array.add(jsonObject.get("maximumGradientNorm"));
        array.add(jsonObject.get("verificationRatio"));
        array.add(jsonObject.get("LogInterval"));
        array.add(jsonObject.get("saveInterval"));
        array.add(jsonObject.get("warmUpSteps"));
        array.add(jsonObject.get("nEFTuneNoiseParameter"));
        array.add(false);
        array.add(false);
        array.add(jsonObject.get("loRaRank"));
        array.add(jsonObject.get("loRaRandomDiscard"));
        array.add("");
        array.add("");
        array.add(true);
        array.add(jsonObject.get("dpoBetaData"));
        array.add(new JSONArray());
        array.add("train_" + runId);
        System.out.println(array);
        return ResponseResult.success(array);
    }

    /**
     * 发布运行任务。
     *
     * @param runPublishDto 更新对象。
     * @return 应答结果对象。
     */
    @OperationLog(type = SysOperationLogType.PUBLISH)
    @PostMapping("/publish")
    public ResponseResult<Void> publish(@MyRequestBody RunPublishDto runPublishDto) {
        String errorMessage = MyCommonUtil.getModelValidationError(runPublishDto, true);
        if (errorMessage != null) {
            return ResponseResult.error(ErrorCodeEnum.DATA_VALIDATED_FAILED, errorMessage);
        }
        if (!this.tuningRunService.publishToModelVersion(runPublishDto)) {
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST);
        }
        return ResponseResult.success();
    }

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

    /**
     * 列出符合过滤条件的精调任务运行列表。
     *
     * @param tuningRunDtoFilter 过滤对象。
     * @param orderParam         排序参数。
     * @param pageParam          分页参数。
     * @return 应答结果对象，包含查询结果集。
     */
    @PostMapping("/list")
    public ResponseResult<MyPageData<TuningRunVo>> list(
            @MyRequestBody TuningRunDto tuningRunDtoFilter,
            @MyRequestBody MyOrderParam orderParam,
            @MyRequestBody MyPageParam pageParam) {
        if (pageParam != null) {
            PageMethod.startPage(pageParam.getPageNum(), pageParam.getPageSize());
        }
        TuningRun tuningRunFilter = MyModelUtil.copyTo(tuningRunDtoFilter, TuningRun.class);
        String orderBy = MyOrderParam.buildOrderBy(orderParam, TuningRun.class);
        List<TuningRun> tuningRunList = tuningRunService.getTuningRunListWithRelation(tuningRunFilter, orderBy);
        return ResponseResult.success(MyPageUtil.makeResponseData(tuningRunList, TuningRun.INSTANCE));
    }

    /**
     * 列出符合过滤条件的SFT任务列表。
     *
     * @return 应答结果对象，包含查询结果集。
     */
    @PostMapping("/listForTree")
    public ResponseResult<List<TuningRunVo>> listForTree(@MyRequestBody TuningRunDto tuningRunDtoFilter) {
        TuningRun tuningRunFilter = MyModelUtil.copyTo(tuningRunDtoFilter, TuningRun.class);
        List<TuningRun> tuningRunList =
                tuningRunService.getTuningRunListWithRelation(tuningRunFilter, "");
        return ResponseResult.success(TuningRun.INSTANCE.fromModelList(tuningRunList));
    }

    /**
     * 提交运行状态信息。
     *
     * @return 应答结果对象，包含查询结果集。
     */
    @NoAuthInterface
    @OperationLog(type = SysOperationLogType.UPLOAD, saveResponse = false)
    @PostMapping("/postRunStatus")
    public ResponseResult<Void> postRunStatus(@MyRequestBody Long runId,
                                              @MyRequestBody Long runTime,
                                              @MyRequestBody Integer runStatus) {
        String errorMessage;
        TuningRun tuningRun = this.tuningRunService.getById(runId);
        if (tuningRun == null) {
            errorMessage = "数据验证失败，当前微调运行并不存在，请联系管理员！";
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST, errorMessage);
        }
        tuningRun.setRunStatus(runStatus);
        tuningRun.setRunTime(runTime);
        tuningRun.setUpdateTime(new Date());
        if (!tuningRunService.updateStatusById(tuningRun)) {
            errorMessage = "运行状态信息提交错误！";
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST, errorMessage);
        }
        return ResponseResult.success();
    }

    /**
     * 提交发布状态信息。
     *
     * @return 应答结果对象，包含查询结果集。
     */
    @NoAuthInterface
    @OperationLog(type = SysOperationLogType.UPLOAD, saveResponse = false)
    @PostMapping("/postPublishStatus")
    public ResponseResult<Void> postPublishStatus(@MyRequestBody Long runId,
                                                  @MyRequestBody Long runTime,
                                                  @MyRequestBody Integer publishStatus) {
        String errorMessage;
        TuningRun tuningRun = this.tuningRunService.getById(runId);
        if (tuningRun == null) {
            errorMessage = "数据验证失败，当前微调运行并不存在，请联系管理员！";
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST, errorMessage);
        }
        tuningRun.setPublishStatus(publishStatus);
        tuningRun.setUpdateTime(new Date());
        if (!tuningRunService.updateById(tuningRun)) {
            errorMessage = "运行发布信息提交错误！";
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST, errorMessage);
        }
        return ResponseResult.success();
    }

    /**
     * 查看指定精调任务运行对象详情。
     *
     * @param runId 指定对象主键Id。
     * @return 应答结果对象，包含对象详情。
     */
    @GetMapping("/view")
    public ResponseResult<TuningRunVo> view(@RequestParam Long runId) {
        TuningRun tuningRun = tuningRunService.getByIdWithRelation(runId, MyRelationParam.full());
        if (tuningRun == null) {
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST);
        }
        TuningRunVo tuningRunVo = TuningRun.INSTANCE.fromModel(tuningRun);
        return ResponseResult.success(tuningRunVo);
    }

    /**
     * 获取运行的日志。
     *
     * @param runId 指定对象主键Id。
     * @return 应答结果对象，包含对象详情。
     */
    @GetMapping("/getLog")
    public ResponseEntity<Resource> getLog(@RequestParam Long runId) throws IOException {
        TuningRun tuningRun = this.tuningRunService.getById(runId);
        ModelVersion modelVersion = this.modelVersionService.getById(tuningRun.getModelVersionId());
        String url = this.pythonConfig.getModelTuningFileBaseDir() + modelVersion.getVersionName() + File.separator + tuningRun.getTrainMethod() + File.separator + "train_" + tuningRun.getRunId() + File.separator + "run_train.log";
        return FileUtil.getFileByUrl(url);
    }

    /**
     * 获取运行状态。
     *
     * @param runId 指定对象主键Id。
     * @return 应答结果对象，包含对象详情。
     */
    @GetMapping("/getStatus")
    public ResponseResult<String> getStatus(@RequestParam Long runId) throws IOException {
        TuningRun tuningRun = this.tuningRunService.getById(runId);
        ModelVersion modelVersion = this.modelVersionService.getById(tuningRun.getModelVersionId());
        String url = this.pythonConfig.getModelTuningFileBaseDir() + modelVersion.getVersionName() + File.separator + tuningRun.getTrainMethod() + File.separator + "train_" + tuningRun.getRunId() + File.separator + "trainer_state.json";
        File file = new File(url); // 指定文件路径
        ObjectMapper objectMapper = new ObjectMapper();
        return ResponseResult.success(objectMapper.readTree(file).toString());
    }


    /**
     * 查看指定精调任务运行新版本对象详情。
     *
     * @param taskId 指定对象主键Id。
     * @return 应答结果对象，包含对象详情。
     */
    @GetMapping("/lastRunVersion")
    public ResponseResult<TuningRunVo> lastRunVersion(@RequestParam Long taskId) {
        TuningRun tuningRunFilter = new TuningRun();
        tuningRunFilter.setTaskId(taskId);
        List<TuningRun> tuningRuns = tuningRunService.getTuningRunList(tuningRunFilter, "run_version");
        if (tuningRuns == null) {
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST);
        }
        TuningRun tuningRun = tuningRuns.get(tuningRuns.size() - 1);
        TuningRunVo tuningRunVo = TuningRun.INSTANCE.fromModel(tuningRun);
        return ResponseResult.success(tuningRunVo);
    }

    private ResponseResult<Void> doDelete(Long runId) {
        String errorMessage;
        if (tuningRunService.getById(runId) == null) {
            errorMessage = "数据验证失败，当前 [对象] 并不存在，请刷新后重试！";
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST, errorMessage);
        }
        if (!tuningRunService.remove(runId)) {
            errorMessage = "数据操作失败，删除的对象不存在，请刷新后重试！";
            return ResponseResult.error(ErrorCodeEnum.DATA_NOT_EXIST, errorMessage);
        }
        return ResponseResult.success();
    }
}
