package com.iplatform.api;

import com.iplatform.base.SystemController;
import com.iplatform.milvus.EventVo;
import com.iplatform.milvus.MilvusEngine;
import com.iplatform.milvus.ParamList;
import com.iplatform.milvus.service.EventServiceImpl;
import com.walker.infrastructure.utils.StringUtils;
import com.walker.web.ResponseValue;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.client.RestTemplate;

import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.concurrent.TimeUnit;

@RestController
@RequestMapping("/test/milvus")
public class MilvusChatApi extends SystemController {

    private EventServiceImpl eventService;
    private RestTemplate restTemplate;
    private MilvusEngine milvusEngine;

    private Long existId = 4048L;
    private boolean isBreak = false;

    private static final String URL_EMBEDDING = "http://120.26.128.84:7003/ai/text/embedding";
    private static final String URL_MILVUS = "120.26.128.84";

    @Autowired
    public MilvusChatApi(EventServiceImpl eventService, RestTemplate restTemplate){
        this.eventService = eventService;
        this.restTemplate = restTemplate;
        if(this.milvusEngine == null){
            MilvusEngine engine = new MilvusEngine(URL_MILVUS, 19530);
            this.milvusEngine = engine;
            logger.info("milvus engine ok!");
        }
    }

    @RequestMapping("/embedding")
    public ResponseValue testHttpEmbedding(){
        List<EventVo> data = new ArrayList(8);
        EventVo eventVo = new EventVo();
        eventVo.setContent("第一句");
        data.add(eventVo);

        eventVo = new EventVo();
        eventVo.setContent("第二句");
        data.add(eventVo);

        boolean success = this.acquireEmbedding(data);
        this.milvusEngine.insertTestData();

        return ResponseValue.success("结果是：" + success);
    }

    @RequestMapping("/write")
    public ResponseValue testWriteMilvus(){
        Collection<EventVo> eventVoList = this.acquireEventVoList();
        if(eventVoList == null){
            return ResponseValue.error("没有加载到数据");
        }
        logger.info("加载了 event vo: {}个", eventVoList.size());

        new Thread(new WriteTask(eventVoList)).start();
        return ResponseValue.success();
    }

    @RequestMapping("/query")
    public ResponseValue testQueryMilvus(String text){

        return ResponseValue.success();
    }

    private boolean acquireEmbedding(List<EventVo> batchData){
        ParamList paramList = new ParamList();
        for(EventVo eventVo : batchData){
            paramList.add(eventVo.getContent());
        }
        try{
            ResponseEntity<ResponseValue> responseEntity = this.restTemplate.postForEntity(URL_EMBEDDING, paramList, ResponseValue.class);
            if(responseEntity.getStatusCode() == HttpStatus.OK){
                List<List<Float>> data = (List<List<Float>>)responseEntity.getBody().getData();
//                double[] one = null;
                List<Float> vector = null;
                for(int i=0; i<data.size(); i++){
                    vector = data.get(i);
//                    one = new double[vector.size()];
//                    for(int j=0; j< vector.size(); j++){
//                        one[j] = vector.get(j);
//                    }

//                    logger.debug("data = {}", vector);
//                    logger.debug("class type = {}", vector.getClass().getName());
                    batchData.get(i).setEmbedding(this.transfer2FloatList(vector));
                }
                return true;
            } else {
                logger.error("http 返回错误：{}", responseEntity.getBody());
            }
            return false;
        } catch (Exception cause){
            logger.error("获取向量出现错误：{}" + cause.getMessage(), cause);
            return false;
        }
    }

    private List<Float> transfer2FloatList(List<?> list){
        List<Float> vector = new ArrayList<>(768);
        for(int i=0; i<list.size(); i++){
            vector.add(Float.parseFloat(list.get(i).toString()));
        }
        return vector;
    }

    private class WriteTask implements Runnable{

        private Collection<EventVo> eventVoList;

        public WriteTask(Collection<EventVo> eventVoList){
            this.eventVoList = eventVoList;
        }

        @Override
        public void run() {
            logger.info(".......... start task ...");

            // 复位状态
            isBreak = false;

            int count = 0;
            List<EventVo> batchData = new ArrayList<>();
            for(EventVo eventVo : this.eventVoList){
                // 已存在上次加载记录，已写入的不在重新处理
                if(existId != null && eventVo.getId() < existId.longValue()){
                    continue;
                }

                if(isBreak){
                    break;
                }

                if(count == 0){
                    logger.info("1) 开始（或继续）采集第 {} 记录", eventVo.getId());
                }

                if(count >= 8){
                    // 触发一次批量写入
                    logger.info("2) 触发一次调用：{}", batchData.get(7).getId());
                    try {
                        boolean successEmbedding = acquireEmbedding(batchData);
                        if(!successEmbedding){
                            logger.error("获取向量失败，任务结束");
                            break;
                        }

                        // 获取向量，写入数据库
                        milvusEngine.insertEventVoList(batchData);

                    } catch (Exception ex){
                        logger.error("error = " + ex.getMessage(), ex);
//                        existId = eventVo.getId();
                        isBreak = true;
                        logger.error("3) 采集任务异常， 当前 id = {}", existId == null ? "" : existId);
                        break;

                    } finally {
                        try {
                            TimeUnit.SECONDS.sleep(1);
                        } catch (InterruptedException e) {
                            e.printStackTrace();
                        }
                    }

                    // 清理本次数据
                    batchData.clear();
                    count = 0;
                }
                batchData.add(eventVo);
                count ++;
            }
            logger.info(".......... end task ...");
        }
    }

    private Collection<EventVo> acquireEventVoList(){
        List<EventVo> data = null;
//        if(this.isBreak){
//            data = this.eventService.queryEventAll(existId);
//        } else {
//            data = this.eventService.queryEventAll(null);
//        }
        data = this.eventService.queryEventAll(null);
        if(StringUtils.isEmptyList(data)){
            return null;
        }

        // 过滤掉 content 重复的数据记录
        LinkedHashMap<String, EventVo> cache = new LinkedHashMap();
        EventVo temp = null;
        for(EventVo e : data){
            temp = cache.get(e.getContent());
            if(temp == null){
                cache.put(e.getContent(), e);
            }
        }
        return cache.values();
    }
}