package org.jk.rag.controller;

import com.alibaba.fastjson2.JSONObject;
import com.baomidou.mybatisplus.core.metadata.IPage;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import jakarta.servlet.http.HttpServletResponse;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.poi.ss.formula.functions.T;
import org.jk.common.core.domain.R;
import org.jk.common.idempotent.annotation.RepeatSubmit;
import org.jk.common.log.annotation.Log;
import org.jk.common.log.enums.BusinessType;
import org.jk.common.mybatis.core.page.PageQuery;
import org.jk.common.mybatis.core.page.TableDataInfo;
import org.jk.common.satoken.utils.LoginHelper;
import org.jk.common.web.core.BaseController;
import org.jk.rag.domain.dto.JkSliceImage;
import org.jk.rag.domain.dto.SliceIdDTO;
import org.jk.rag.domain.entity.*;
import org.jk.rag.mapper.JkDocumentMapper;
import org.jk.rag.mapper.JkSliceInfoMapper;
import org.jk.rag.service.IJkRagService;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

@Slf4j
@RestController
@RequiredArgsConstructor
@RequestMapping("/deepseek/api")
public class JkRagSliceController extends BaseController {


    private final IJkRagService iJkRagService;

    private final JkDocumentMapper jkDocumentMapper;

    private final JkSliceInfoMapper jkSliceInfoMapper;


    /**
     * 上传切片图片
     *
     * @param parms
     * @param knowledgeId
     * @return
     * @throws Exception
     */
    @PostMapping("/uploadSliceImage/{knowledgeId}/{documentId}")
    public R<JkMediaReplacement> uploadSliceImage(@RequestBody JkSliceImpageParms parms,
                                     @PathVariable String knowledgeId,
                                     @PathVariable String documentId) throws Exception {
        try {
            JkMediaReplacement jkMediaReplacement = iJkRagService.uploadSliceImage(parms, knowledgeId, documentId);
            return R.ok(jkMediaReplacement);
        } catch (Exception e) {
            return R.fail(e.getMessage());
        }
    }


    /**
     * 获取切片列表
     *
     * @param params
     * @return
     */
    @PostMapping("/getSliceList")
    public TableDataInfo<JkSliceInfo> getDialogDetail(@RequestBody JkSliceParams params,@RequestBody PageQuery pageQuery) {
        Page<JkSliceInfo> sliceList = iJkRagService.selectPageByDocumentId(params.getDocument_id(), params.getText(),pageQuery);
        return TableDataInfo.build(sliceList);
    }

    /**
     * 根据文档获取切片列表
     *
     * @param documentId
     * @return
     */
    @GetMapping("/getSliceListByDocumentId/{documentId}")
    public R<List<JkSliceInfo>> getSliceListByDocumentId(@PathVariable String documentId) {
        List<JkSliceInfo> sliceList = iJkRagService.selectByDocumentId(documentId, null);
        return R.ok(sliceList);
    }

    /**
     * 删除切片
     *
     * @param sliceId
     * @param knowledgeId
     */
    @DeleteMapping("/deleteSlice/{sliceId}/{knowledgeId}/{documentId}")
    public R<Integer> deleteSlice(@PathVariable String sliceId, @PathVariable String knowledgeId, @PathVariable String documentId) {
        Integer result = iJkRagService.deleteSlice(sliceId, knowledgeId, documentId);
        if(result == 1) {
            iJkRagService.updateDocumentKnowledgeSize(knowledgeId, documentId);
        }
        return R.ok(result);
    }

    /**
     * 切片详情
     *
     * @param sliceId
     * @param knowledgeId
     * @return
     */
    @GetMapping("/getSliceDetail/{sliceId}/{knowledgeId}")
    public R<JSONObject> test(@PathVariable String sliceId, @PathVariable String knowledgeId) {
        JSONObject json = iJkRagService.getSliceDetail(sliceId, knowledgeId);
        return R.ok(json);
    }

    /**
     * 更新切片信息
     *
     * @param params
     * @return
     */
    @RepeatSubmit
    @PutMapping("/updateSliceInfo")
    public R<Integer> updateSliceInfo(@RequestBody JkSliceUpdateParams params) {
        if (params.getSliceIndex().intValue() != params.getIndex().intValue()) {
            jkSliceInfoMapper.updateSliceIndex(params.getIndex(), params.getSliceId());
        }
        Integer result = iJkRagService.updateSliceInfo(params,true);
        if(result == 1) {
            iJkRagService.updateDocumentKnowledgeSize(params.getKnowledgeId(), params.getDocumentId());
        }
        return R.ok(result);
    }


    /**
     * 新增切片
     *
     * @param params
     * @return
     */
    @RepeatSubmit
    @PostMapping("/add/slice")
    public R<String> addSlice(@RequestBody JkSliceUpdateParams params) {
        String result = iJkRagService.addSlice(params);
        if(result.equals("1")) {
            iJkRagService.updateDocumentKnowledgeSize(params.getKnowledgeId(), params.getDocumentId());
        }
        return R.ok(result);
    }

    /**
     * 获取聊天切片列表
     *
     * @param params
     * @return
     */
    @PostMapping("/slicePage")
    public TableDataInfo<SliceInfo> dataExportList(@RequestBody SlicePageInfo params) {
        DocumentInfo chunkInfo = params.getChunk_info();
        List<ChunkList> chunkList = chunkInfo.getChunk_list();
        // 1. 先按 rerank_score 降序排序所有chunk
        List<ChunkList> sortedChunks = chunkList.stream()
                .sorted(Comparator.comparingDouble(
                        (ChunkList s) -> Double.parseDouble(s.getRerank_score())
                ).reversed())
                .collect(Collectors.toList());

        // 2. 提取排序后的chunk_ids（此时已按分数排序）
        List<String> sortedChunkIds = sortedChunks.stream()
                .map(ChunkList::getChunk_id)
                .collect(Collectors.toList());

        int fromIndex = (params.getPageNum() - 1) * params.getPageSize();

        int toIndex = Math.min(fromIndex + params.getPageSize(), chunkList.size());

        List<String> ids = sortedChunkIds.subList(fromIndex, toIndex);

        //PageHelper.startPage(params.getPageNum(), params.getPageSize(), "").setReasonable(true);
        List<SliceInfo> list = iJkRagService.slicePageInfoList(ids);
        list.forEach(sliceInfo -> {
            chunkList.forEach(chunk -> {
                if (sliceInfo.getSliceId().equals(chunk.getChunk_id())) {
                    sliceInfo.setRerankScore(chunk.getRerank_score());
                    sliceInfo.setChunkLength(chunk.getChunk_len());
                }
            });
        });
        List<SliceInfo> sliceList = list.stream()
                .sorted(Comparator.comparingDouble((SliceInfo s) -> Double.parseDouble(s.getRerankScore())).reversed())
                .collect(Collectors.toList());
        return TableDataInfo.build(sliceList);
    }

    /**
     * 获取聊天切片列表
     *
     * @param params
     * @return
     */
    @PostMapping("/free/slicePage")
    public TableDataInfo freeDataExportList(@RequestBody SlicePageInfo params) {
        return dataExportList(params);
    }

    @GetMapping("/slice/search/{id}")
    public R<JSONObject> searchSlice(@PathVariable String id) {
        return R.ok(iJkRagService.searchSlice(id));
    }

    @GetMapping("/free/slice/search/{id}")
    public R<JSONObject> freeSearchSlice(@PathVariable String id) {
        return searchSlice(id);
    }

    /**
     * 切片文档列表
     *
     * @param params
     * @return
     */
    @PostMapping("/sliceDocumentList")
    public TableDataInfo sliceDocumentList(@RequestBody JkDocumentParams params,@RequestBody PageQuery pageQuery) {
        IPage<JkSliceDocument> list = iJkRagService.sliceDocumentPageList(params,pageQuery);
        return TableDataInfo.build(list);
    }

    @Log(title = "下载PDF文档", businessType = BusinessType.EXPORT)
    @PostMapping("/pdfStream")
    public void  pdfStream(HttpServletResponse response, @RequestBody JkDocumentExportParams parms) throws IOException {

        // 设置响应头
        response.setContentType("application/octet-stream");
        response.setCharacterEncoding("UTF-8");
        response.setHeader("Content-Disposition", "attachment; filename=\"documents.zip\"");

        // 使用 StreamingResponseBody 来实现流式响应
        iJkRagService.pdfStream(parms.getDocumentIds() ,response.getOutputStream());

    }

    /**
     * 更新切片信息
     *
     * @param params
     * @return
     */
    @RepeatSubmit
    @PutMapping("/updateSliceInfoList")
    public R<Map<String,String>> updateSliceInfoList(@RequestBody JkSliceUpdateReq params) {
        Map<String,String> map = new HashMap<>();
        map.put("documentId", params.getDocumentId());
        map.put("knowledgeId", params.getKnowledgeId());
        try{
            BmDocument bmDocument = jkDocumentMapper.selectById(params.getDocumentId());
            Boolean flg = true;
            if(bmDocument.getQaChecked() != params.getQaChecked()
                || bmDocument.getRelatedQuestionsEnabled() != params.getRelatedQuestionsEnabled()
                || bmDocument.getSummaryGenerationEnabled() != params.getSummaryGenerationEnabled()) {
                flg = false;
            }
            iJkRagService.updateDocument(JkDocumentParams.builder()
                .qaChecked(params.getQaChecked())
                .parentGenerationEnabled(params.getParentGenerationEnabled())
                .summaryGenerationEnabled(params.getSummaryGenerationEnabled())
                .relatedQuestionsEnabled(params.getRelatedQuestionsEnabled())
                .status("4")
                .build(),params.getDocumentId());

            //iJkRagService.updateSliceInfoList(params,flg);
            if(null != params.getDelSliceIds()) {
                for(String sliceId : params.getDelSliceIds()) {
                    iJkRagService.deleteSlice(sliceId,params.getKnowledgeId(), params.getDocumentId());
                }
            }
            for(JkSliceUpdateParams slice:params.getSliceList()) {
                if (null == slice.getSliceId()) {
                    slice.setKnowledgeId(params.getKnowledgeId());
                    slice.setDocumentId(params.getDocumentId());
                    iJkRagService.addSlice(slice);
                    continue;
                }
                JkSliceInfo jkSliceInfo = jkSliceInfoMapper.selectBySliceId(slice.getSliceId());
                if (slice.getSliceIndex().intValue() != slice.getIndex().intValue() ||jkSliceInfo.getDelFlag().equals("1")) {
                    jkSliceInfoMapper.updateSliceIndex(slice.getIndex(), slice.getSliceId());
                }
                if (!flg || !slice.getSliceText().equals(jkSliceInfo.getSliceText())) {
                    iJkRagService.updateSliceInfo(slice,flg);
                }
            }
            iJkRagService.updateDocument(JkDocumentParams.builder()
                .status("5")
                .build(),params.getDocumentId());

            return R.ok(map);
        } catch (Exception e) {
            iJkRagService.updateDocument(JkDocumentParams.builder()
                .status("6")
                .build(),params.getDocumentId());
            map.put("errorMsg",e.getMessage());
            return R.fail(map);
        } finally {
            try {
                iJkRagService.updateDocumentKnowledgeSize(params.getKnowledgeId(), params.getDocumentId());
            }catch (Exception e) {
                log.info("更新知识库文档size失败："+e.getMessage());
            }

        }


    }

    /**
     * 根据文档获取初版切片列表
     *
     * @param documentId
     * @return
     */
    @GetMapping("/getOldSliceListByDocumentId/{documentId}")
    public R<List<JkSliceInfo>> getOldSliceListByDocumentId(@PathVariable String documentId) {
        List<JkSliceInfo> sliceList = iJkRagService.selectOldSliceByDocumentId(documentId, null);
        return R.ok(sliceList);
    }

    /**
     * 根据文档获取初版切片列表
     *
     * @param documentId
     * @return
     */
    @GetMapping("/getMarkDodnSliceByDocumentId/{documentId}")
    public R<List<String>> getMarkDodnSliceByDocumentId(@PathVariable String documentId) throws IOException {
        String re = iJkRagService.getMarkDownSliceByDocumentId(documentId);
        return R.ok(List.of(re));
    }

    @PutMapping("/editMarkDownFile")
    public R<String> editMarkDownFile(@RequestBody JkMarkUpdateReq req){
        try{
            iJkRagService.updateDocument(JkDocumentParams.builder()
                .status("4")
                .qaChecked(false)
                .parentGenerationEnabled(false)
                .summaryGenerationEnabled(false)
                .relatedQuestionsEnabled(false)
                .build(),req.getDocumentId());
            iJkRagService.updateDocumentSetting(null, req.getDocumentId());
            iJkRagService.editMarkDownFile(req);
            iJkRagService.updateDocument(JkDocumentParams.builder()
                .status("5")
                .build(),req.getDocumentId());
            return R.ok();
        } catch (Exception e) {
            iJkRagService.updateDocument(JkDocumentParams.builder()
                .status("6")
                .build(),req.getDocumentId());
            return R.fail(e.getMessage());
        }

    }

    /**
     * 删除切片
     *
     * @param parm
     * @param knowledgeId
     */
    @DeleteMapping("/deleteSliceList/{knowledgeId}/{documentId}")
    public R<Integer> deleteSliceList(@RequestBody SliceIdDTO parm, @PathVariable String knowledgeId, @PathVariable String documentId) {
        for(String sliceId :parm.getSliceIds()) {
            int result = iJkRagService.deleteSlice(sliceId, knowledgeId, documentId);
            if(result == 0) {
                iJkRagService.updateDocumentKnowledgeSize(knowledgeId, documentId);
                return R.ok(result);
            }
        }
        return R.ok(1);
    }

}
