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.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.MyRelationParam;
import com.yice.common.core.object.TokenData;
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.ModelCompressMapper;
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("modelCompressService")
public class ModelCompressServiceImpl extends BaseService<ModelCompress, Long> implements ModelCompressService {

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


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

    /**
     * 保存新增对象。
     *
     * @param modelCompress 新增对象。
     * @return 返回新增对象。
     */
    @Transactional(rollbackFor = Exception.class)
    @Override
    public ModelCompress saveNew(ModelCompress modelCompress) {
        ModelVersion modelVersionSource = this.modelVersionService.getById(modelCompress.getSourceVersionId());
        ModelTask modelTask = new ModelTask();
        ModelVersion modelVersionTarget = saveAll(modelCompress, modelVersionSource, modelTask);
        messageWithSocket(modelCompress, modelVersionSource, modelVersionTarget, modelTask);
        return modelCompress;
    }

    @SneakyThrows
    private void messageWithSocket(ModelCompress modelCompress, ModelVersion sourceModelVersion, 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);
                String receiveMsg = receiveJson.getString("msg");
                JSONObject sendJson = new JSONObject();
                if (receiveMsg.equals("send_hash")) {
                    sendJson.put("fn_index", 44);
                    sendJson.put("session_hash", String.valueOf(modelCompress.getTaskId()));
                    log.info("发送服务端的消息:" + sendJson.toJSONString());
                    this.send(sendJson.toJSONString());
                } else if (receiveMsg.equals("send_data")) {
                    JSONArray array = new JSONArray();
                    array.add("zh");
                    array.add(sourceModelVersion.getVersionName());
                    array.add(sourceModelVersion.getModelUrl());
                    array.add(new JSONArray());
                    array.add("lora");
                    array.add(sourceModelVersion.getBasePromptTemplate());
                    array.add(2);
                    array.add(pythonConfig.getModelOutputFileBaseDir() + modelVersionTarget.getVersionName());
                    array.add("8");
                    sendJson.put("data", array);
                    sendJson.put("event_data", "null");
                    sendJson.put("fn_index", 44);
                    sendJson.put("session_hash", String.valueOf(modelCompress.getTaskId()));
                    log.info("发送服务端的消息:" + sendJson.toJSONString());
                    this.send(sendJson.toJSONString());
                } else if (receiveMsg.equals("process_completed")) {
                    updateStatus(receiveJson.getBoolean("success"), modelCompress, 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(ModelCompress modelCompress, ModelVersion modelVersionSource, ModelTask modelTask) {
        modelCompress.setTaskStatus(3);
        modelCompressMapper.insert(this.buildDefaultValue(modelCompress));
        ModelVersion modelVersion = new ModelVersion();
        modelVersion.setTaskId(modelCompress.getTaskId());
        modelVersion.setBaseModel(modelVersionSource.getBaseModel());
        modelVersion.setBaseId(modelVersionSource.getBaseId());
        modelVersion.setIsCompress(1);
        modelVersion.setStatus(0);
        modelVersion.setBasePromptTemplate(modelVersionSource.getBasePromptTemplate());
        modelVersion.setRunName(modelVersionSource.getRunName());
        modelVersion.setModelTrainingMethod(modelVersionSource.getModelTrainingMethod());
        TuningTask tuningTask = this.tuningTaskService.getById(modelVersionSource.getTaskId());
        modelVersion.setTrainingTask(tuningTask.getTaskName());
        if (modelCompress.getCreateMethod() == 0) {
            modelVersion.setModelId(modelCompress.getTargetModelId());
            modelVersionService.saveNew(modelVersion);
        } else {
            ModelManage modelManage = new ModelManage();
            modelManage.setModelName(modelCompress.getTargetModelName());
            modelManage.setModelDescribe(modelCompress.getTaskDescribe());
            modelManage.setModelType(0);
            modelManage.setIsBaseModel(0);
            modelManageService.saveAndCreateVersion(modelManage, modelVersion);
        }
        modelCompress.setTargetVersionId(modelVersion.getVersionId());
        this.updateById(modelCompress);
        modelTask.setModelVersion(modelVersionSource.getModelVersion());
        modelTask.setModelId(modelVersionSource.getModelId());
        modelTask.setTaskType(2);
        modelTask.setTaskStatus(3);
        modelTask.setVersionId(modelVersionSource.getVersionId());
        modelTask.setVersionName(modelVersionSource.getVersionName());
        modelTask.setTaskId(modelCompress.getTaskId());
        this.modelTaskService.saveNew(modelTask);

        return modelVersion;
    }

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

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

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

    @Transactional(rollbackFor = Exception.class)
    @Override
    public boolean updateWithRelation(
            ModelCompress modelCompress, ModelCompress originalModelCompress, JSONObject relationData) {
        // modelCompress 为空的时候，无需修改主表数据。
        if (modelCompress != null && !this.update(modelCompress, originalModelCompress)) {
            return false;
        }
        this.saveOrUpdateOneToOneRelationData(originalModelCompress, relationData, TokenData.takeFromRequest().getUserId());
        return true;
    }

    private void saveOrUpdateOneToOneRelationData(ModelCompress modelCompress, JSONObject relationData, Long userId) {
        // 对于一对一新增或更新，如果主键值为空就新增，否则就更新，同时更新updateTime和updateUserId。
        ModelTask modelTask = relationData.getObject("modelTask", ModelTask.class);
        if (modelTask != null) {
            modelTask.setTaskId(modelCompress.getTaskId());
            modelTask.setCreateUserId(userId);
            modelTask.setUpdateUserId(userId);
            modelTaskService.saveNew(modelTask);
            /*modelTaskService.saveNewOrUpdate(modelTask,
                    modelTaskService::saveNew, modelTaskService::update);*/
        }
    }

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

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

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


    private ModelCompress buildDefaultValue(ModelCompress modelCompress) {
        if (modelCompress.getTaskId() == null) {
            modelCompress.setTaskId(idGenerator.nextLongId());
        }
        MyModelUtil.fillCommonsForInsert(modelCompress);
        return modelCompress;
    }
}
