package com.yice.webadmin.app.service.impl;

import cn.hutool.core.collection.CollUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.core.conditions.update.UpdateWrapper;
import com.github.pagehelper.Page;
import com.yice.common.core.base.dao.BaseDaoMapper;
import com.yice.common.core.base.service.BaseService;
import com.yice.common.core.object.CallResult;
import com.yice.common.core.object.MyRelationParam;
import com.yice.common.core.util.MyModelUtil;
import com.yice.common.sequence.wrapper.IdGeneratorWrapper;
import com.yice.webadmin.app.config.PythonConfig;
import com.yice.webadmin.app.dao.TuningRunMapper;
import com.yice.webadmin.app.dto.RunPublishDto;
import com.yice.webadmin.app.model.*;
import com.yice.webadmin.app.service.*;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.net.URI;
import java.util.Date;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 精调任务运行数据操作服务类。
 *
 * @author linking
 * @date 2023-04-13
 */
@Slf4j
@Service("tuningRunService")
public class TuningRunServiceImpl extends BaseService<TuningRun, Long> implements TuningRunService {

    @Autowired
    private TuningRunMapper tuningRunMapper;
    @Autowired
    private TuningTaskService tuningTaskService;
    @Autowired
    private ModelManageService modelManageService;
    @Autowired
    private ModelVersionService modelVersionService;
    @Autowired
    private IdGeneratorWrapper idGenerator;
    @Autowired
    private PythonConfig pythonConfig;
    @Autowired
    private ModelTaskService modelTaskService;


    /**
     * 返回当前Service的主表Mapper对象。
     *
     * @return 主表Mapper对象。
     */
    @Override
    protected BaseDaoMapper<TuningRun> mapper() {
        return tuningRunMapper;
    }

    /**
     * 保存新增对象。
     *
     * @param tuningRun 新增对象。
     * @return 返回新增对象。
     */
    @Transactional(rollbackFor = Exception.class)
    @Override
    public TuningRun saveNew(TuningRun tuningRun) {
        TuningRun tuningRunFilter = new TuningRun();
        tuningRunFilter.setTaskId(tuningRun.getTaskId());
        List<TuningRun> reTuningRunList = this.getTuningRunList(tuningRunFilter, "run_version");
        TuningTask tuningTask = this.tuningTaskService.getById(tuningRun.getTaskId());
        Integer version = 1;
        if (reTuningRunList != null && reTuningRunList.size() != 0) {
            version = reTuningRunList.get(reTuningRunList.size() - 1).getRunVersion() + 1;
        }
        tuningRun.setRunVersion(version);
        tuningRun.setRunName(tuningTask.getTaskName() + "_V" + tuningRun.getRunVersion());
        tuningRunMapper.insert(this.buildDefaultValue(tuningRun));
        return tuningRun;
    }

    /**
     * 利用数据库的insertList语法，批量插入对象列表。
     *
     * @param tuningRunList 新增对象列表。
     */
    @Transactional(rollbackFor = Exception.class)
    @Override
    public void saveNewBatch(List<TuningRun> tuningRunList) {
        if (CollUtil.isNotEmpty(tuningRunList)) {
            tuningRunList.forEach(this::buildDefaultValue);
            tuningRunMapper.insertList(tuningRunList);
        }
    }

    /**
     * 更新数据对象。
     *
     * @param tuningRun         更新的对象。
     * @param originalTuningRun 原有数据对象。
     * @return 成功返回true，否则false。
     */
    @Transactional(rollbackFor = Exception.class)
    @Override
    public boolean update(TuningRun tuningRun, TuningRun originalTuningRun) {
        MyModelUtil.fillCommonsForUpdate(tuningRun, originalTuningRun);
        // 这里重点提示，在执行主表数据更新之前，如果有哪些字段不支持修改操作，请用原有数据对象字段替换当前数据字段。
        UpdateWrapper<TuningRun> uw = this.createUpdateQueryForNullValue(tuningRun, tuningRun.getRunId());
        return tuningRunMapper.update(tuningRun, uw) == 1;
    }

    /**
     * 删除指定数据。
     *
     * @param runId 主键Id。
     * @return 成功返回true，否则false。
     */
    @Transactional(rollbackFor = Exception.class)
    @Override
    public boolean remove(Long runId) {
        return tuningRunMapper.deleteById(runId) == 1;
    }

    /**
     * 当前服务的支持表为从表，根据主表的关联Id，删除一对多的从表数据。
     *
     * @param taskId 从表关联字段。
     * @return 删除数量。
     */
    @Transactional(rollbackFor = Exception.class)
    @Override
    public int removeByTaskId(Long taskId) {
        TuningRun deletedObject = new TuningRun();
        deletedObject.setTaskId(taskId);
        return tuningRunMapper.delete(new QueryWrapper<>(deletedObject));
    }

    /**
     * 获取单表查询结果。由于没有关联数据查询，因此在仅仅获取单表数据的场景下，效率更高。
     * 如果需要同时获取关联数据，请移步(getTuningRunListWithRelation)方法。
     *
     * @param filter  过滤对象。
     * @param orderBy 排序参数。
     * @return 查询结果集。
     */
    @Override
    public List<TuningRun> getTuningRunList(TuningRun filter, String orderBy) {
        return tuningRunMapper.getTuningRunList(filter, orderBy);
    }

    /**
     * 获取主表的查询结果，以及主表关联的字典数据和一对一从表数据，以及一对一从表的字典数据。
     * 该查询会涉及到一对一从表的关联过滤，或一对多从表的嵌套关联过滤，因此性能不如单表过滤。
     * 如果仅仅需要获取主表数据，请移步(getTuningRunList)，以便获取更好的查询性能。
     *
     * @param filter  主表过滤对象。
     * @param orderBy 排序参数。
     * @return 查询结果集。
     */
    @Override
    public List<TuningRun> getTuningRunListWithRelation(TuningRun filter, String orderBy) {
        List<TuningRun> resultList = tuningRunMapper.getTuningRunList(filter, orderBy);
        // 在缺省生成的代码中，如果查询结果resultList不是Page对象，说明没有分页，那么就很可能是数据导出接口调用了当前方法。
        // 为了避免一次性的大量数据关联，规避因此而造成的系统运行性能冲击，这里手动进行了分批次读取，开发者可按需修改该值。
        int batchSize = resultList instanceof Page ? 0 : 1000;
        this.buildRelationForDataList(resultList, MyRelationParam.normal(), batchSize);
        return resultList;
    }


    @Override
    @Transactional(rollbackFor = Exception.class)
    public boolean publishToModelVersion(RunPublishDto runPublishDto) {
        TuningRun tuningRun = this.getById(runPublishDto.getRunId());
        ModelVersion modelVersionSource = this.modelVersionService.getById(tuningRun.getModelVersionId());
        ModelTask modelTask = new ModelTask();
        ModelVersion modelVersionTarget = saveAll(tuningRun, runPublishDto, modelVersionSource, modelTask);
        messageWithSocket(tuningRun, modelVersionSource, modelVersionTarget, modelTask);
        return true;
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public ModelManage createToModel(ModelManage modelManage, ModelVersion modelVersion) {
        TuningRun tuningRun = this.getById(modelVersion.getRunId());
        ModelVersion modelVersionSource = this.modelVersionService.getById(tuningRun.getModelVersionId());
        ModelTask modelTask = new ModelTask();
        ModelVersion modelVersionTarget = saveAll(tuningRun, modelManage, modelVersion, modelVersionSource, modelTask);
        messageWithSocket(tuningRun, modelVersionSource, modelVersionTarget, modelTask);
        return modelManage;
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public ModelVersion createToModelVersion(ModelVersion modelVersion) {
        TuningRun tuningRun = this.getById(modelVersion.getRunId());
        ModelVersion modelVersionSource = this.modelVersionService.getById(tuningRun.getModelVersionId());
        ModelTask modelTask = new ModelTask();
        ModelVersion modelVersionTarget = saveAll(tuningRun, modelVersion, modelVersionSource, modelTask);
        messageWithSocket(tuningRun, modelVersionSource, modelVersionTarget, modelTask);
        return modelVersionTarget;
    }

    @Override
    public boolean updateStatusById(TuningRun tuningRun) {
        return this.updateById(tuningRun);
    }


    @SneakyThrows
    private void messageWithSocket(TuningRun tuningRun, ModelVersion modelVersionSource, ModelVersion modelVersionTarget, ModelTask modelTask) {
        new WebSocketClient(new URI(this.pythonConfig.getPythonWebsocketUri())) {
            @Override
            public void onOpen(ServerHandshake serverHandshake) {
                log.info("-------------与大模型建立连接-------------");
            }

            @Override
            public void onMessage(String message) {
                log.info("收到来自服务端的消息:" + message);
                JSONObject receiveJson = (JSONObject) JSON.parse(message);
                JSONObject jsonObject = (JSONObject) JSON.parse(tuningRun.getConfiguration());
                String receiveMsg = receiveJson.getString("msg");
                JSONObject sendJson = new JSONObject();
                if (receiveMsg.equals("send_hash")) {
                    System.out.println("isSuccess1：" + System.currentTimeMillis());
                    sendJson.put("fn_index", 44);
                    sendJson.put("session_hash", tuningRun.getRunId().toString());
                    this.send(sendJson.toJSONString());
                }
                if (receiveMsg.equals("send_data")) {
                    System.out.println("isSuccess2：" + System.currentTimeMillis());
                    JSONArray array = new JSONArray();
                    array.add("zh");
                    array.add(modelVersionSource.getVersionName());
                    array.add(modelVersionSource.getModelUrl());
                    JSONArray ddJson = new JSONArray();
                    ddJson.add("train_" + tuningRun.getRunId());
                    array.add(ddJson);
                    array.add(tuningRun.getTrainMethod());
                    array.add(jsonObject.get("promptTemplate"));
                    array.add(2);
                    array.add(pythonConfig.getModelOutputFileBaseDir() + modelVersionTarget.getVersionName());
                    array.add("none");
                    sendJson.put("data", array);
                    sendJson.put("event_data", "null");
                    sendJson.put("fn_index", 44);
                    sendJson.put("session_hash", tuningRun.getRunId().toString());
                    this.send(sendJson.toJSONString());
                }
                log.info("发送服务端的消息:" + sendJson.toJSONString());
                if (receiveMsg.equals("process_completed")) {
                    System.out.println("isSuccess4：" + System.currentTimeMillis());
                    updateStatus(receiveJson.getBoolean("success"), tuningRun, modelVersionTarget, modelTask);
                }
            }

            @Override
            public void onClose(int i, String s, boolean b) {
                log.info("关闭连接:::" + "i = " + i + ":::s = " + s + ":::b = " + b);
            }

            @Override
            public void onError(Exception e) {
                log.error("报错了:::" + e.getMessage());
            }
        }.connectBlocking(5000, TimeUnit.MILLISECONDS);
    }

    public ModelVersion saveAll(TuningRun tuningRun, RunPublishDto runPublishDto, ModelVersion modelVersionSource, ModelTask modelTask) {
        ModelVersion modelVersion = new ModelVersion();
        this.initModelVersion(modelVersion, tuningRun, modelVersionSource);
        ModelManage modelManage = new ModelManage();
        if (runPublishDto.getPublishWay() == 0) {
            modelManage.setModelDescribe(runPublishDto.getModelDescribe());
            modelManage.setModelName(runPublishDto.getModelName());
            modelManage.setModelType(0);
            modelManage.setIsBaseModel(0);
            modelManageService.saveAndCreateVersionV(modelManage, modelVersion);
        } else {
            modelVersion.setModelId(runPublishDto.getModelId());
            modelVersionService.saveNew(modelVersion);
        }
        this.initModelTask(modelTask, tuningRun, modelVersion);
        this.modelTaskService.saveNew(modelTask);
        return modelVersion;
    }

    public ModelVersion saveAll(TuningRun tuningRun, ModelManage modelManageS, ModelVersion modelVersion, ModelVersion modelVersionSource, ModelTask modelTask) {
        this.initModelVersion(modelVersion, tuningRun, modelVersionSource);
        modelManageS.setModelType(0);
        modelManageS.setIsBaseModel(0);
        modelManageService.saveAndCreateVersionV(modelManageS, modelVersion);
        this.initModelTask(modelTask, tuningRun, modelVersion);
        this.modelTaskService.saveNew(modelTask);
        return modelVersion;
    }


    public ModelVersion saveAll(TuningRun tuningRun, ModelVersion modelVersion, ModelVersion modelVersionSource, ModelTask modelTask) {
        this.initModelVersion(modelVersion, tuningRun, modelVersionSource);
        modelVersionService.saveNew(modelVersion);
        modelTask.setModelVersion(modelVersion.getModelVersion());
        this.initModelTask(modelTask, tuningRun, modelVersion);
        this.modelTaskService.saveNew(modelTask);
        return modelVersion;
    }

    private void initModelVersion(ModelVersion modelVersion, TuningRun tuningRun, ModelVersion modelVersionSource) {
        modelVersion.setRunId(tuningRun.getRunId());
        modelVersion.setTaskId(tuningRun.getTaskId());
        modelVersion.setBaseModel(modelVersionSource.getVersionName());
        modelVersion.setBaseId(modelVersionSource.getVersionId());
        modelVersion.setIsCompress(0);
        modelVersion.setStatus(0);
        modelVersion.setBasePromptTemplate(modelVersionSource.getBasePromptTemplate());
        modelVersion.setRunName(tuningRun.getRunName());
        modelVersion.setModelTrainingMethod(tuningRun.getTrainMethod());
        TuningTask tuningTask = this.tuningTaskService.getById(tuningRun.getTaskId());
        modelVersion.setTrainingTask(tuningTask.getTaskName());
    }

    private void initModelTask(ModelTask modelTask, TuningRun tuningRun, ModelVersion modelVersion) {
        modelTask.setModelVersion(modelVersion.getModelVersion());
        modelTask.setModelId(modelVersion.getModelId());
        modelTask.setTaskType(0);
        modelTask.setTaskStatus(3);
        modelTask.setVersionId(modelVersion.getVersionId());
        modelTask.setVersionName(modelVersion.getVersionName());
        modelTask.setTaskId(tuningRun.getTaskId());
    }

    public Boolean updateStatus(Boolean flag, TuningRun tuningRun, ModelVersion modelVersionTarget, ModelTask modelTask) {
        if (flag) {
            modelVersionTarget.setModelUrl(pythonConfig.getModelOutputFileBaseDir() + modelVersionTarget.getVersionName());
            modelVersionTarget.setStatus(1);
            modelTask.setTaskStatus(1);
            tuningRun.setPublishStatus(1);
        } else {
            modelTask.setTaskStatus(-1);
            modelVersionTarget.setStatus(-1);
            tuningRun.setPublishStatus(-1);
        }
        modelTask.setCompleteTime(new Date());
        modelTaskService.updateById(modelTask);
        modelVersionService.updateById(modelVersionTarget);
        return updateById(tuningRun);
    }


    /**
     * 根据最新对象和原有对象的数据对比，判断关联的字典数据和多对一主表数据是否都是合法数据。
     *
     * @param tuningRun         最新数据对象。
     * @param originalTuningRun 原有数据对象。
     * @return 数据全部正确返回true，否则false。
     */
    @Override
    public CallResult verifyRelatedData(TuningRun tuningRun, TuningRun originalTuningRun) {
        String errorMessageFormat = "数据验证失败，关联的%s并不存在，请刷新后重试！";
        //这里是一对多的验证
        if (this.needToVerify(tuningRun, originalTuningRun, TuningRun::getTaskId)
                && !tuningTaskService.existId(tuningRun.getTaskId())) {
            return CallResult.error(String.format(errorMessageFormat, "任务ID"));
        }
        return CallResult.ok();
    }

    private TuningRun buildDefaultValue(TuningRun tuningRun) {
        if (tuningRun.getRunId() == null) {
            tuningRun.setRunId(idGenerator.nextLongId());
        }
        MyModelUtil.fillCommonsForInsert(tuningRun);
        return tuningRun;
    }
}
