/*
 * Copyright (c) 2024 LangChat. TyCoding All Rights Reserved.
 *
 * Licensed under the GNU Affero General Public License, Version 3 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.gnu.org/licenses/agpl-3.0.html
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package cn.iocoder.yudao.module.langchat.controller.admin.endpoint;

import ch.qos.logback.core.net.LoginAuthenticator;
import cn.hutool.core.util.StrUtil;

import cn.iocoder.yudao.framework.common.exception.ServiceException;
import cn.iocoder.yudao.framework.common.pojo.CommonResult;
import cn.iocoder.yudao.framework.tenant.core.aop.TenantIgnore;
import cn.iocoder.yudao.module.langchat.dal.dto.ChatReq;
import cn.iocoder.yudao.module.langchat.dal.dto.EmbeddingR;
import cn.iocoder.yudao.module.langchat.entity.AigcDocs;
import cn.iocoder.yudao.module.langchat.entity.AigcDocsSlice;
import cn.iocoder.yudao.module.langchat.entity.AigcOss;
import cn.iocoder.yudao.module.langchat.enums.EmbedConst;
import cn.iocoder.yudao.module.langchat.mapper.AigcDocsMapper;
import cn.iocoder.yudao.module.langchat.service.aigc.AigcKnowledgeService;
import cn.iocoder.yudao.module.langchat.service.aigc.AigcOssService;
import cn.iocoder.yudao.module.langchat.service.core.EmbeddingService;
import cn.iocoder.yudao.module.langchat.service.core.LangEmbeddingService;
import cn.iocoder.yudao.module.langchat.task.TaskManager;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.util.concurrent.Executors;

import static cn.iocoder.yudao.framework.security.core.util.SecurityFrameworkUtils.getLoginUserId;

/**
 * @author tycoding
 * @since 2024/4/25
 */
@Slf4j
@RestController
@AllArgsConstructor
@RequestMapping("/chat/aigc/embedding")
public class EmbeddingEndpoint {

    private final LangEmbeddingService langEmbeddingService;
    private final AigcKnowledgeService aigcKnowledgeService;
    private final AigcDocsMapper aigcDocsMapper;
    private final AigcOssService aigcOssService;
    private final EmbeddingService embeddingService;

    @PostMapping("/text")
    @PreAuthorize("@ss.hasPermission('aigc:embedding:text')")
    @TenantIgnore
    public CommonResult text(@RequestBody AigcDocs data) {
        if (StrUtil.isBlankIfStr(data.getContent())) {
            throw new ServiceException(-1003,"文档内容不能为空");
        }
        if (StrUtil.isBlank(data.getId())) {
            aigcKnowledgeService.addDocs(data);
        }
        data.setType(EmbedConst.ORIGIN_TYPE_INPUT).setSliceStatus(false);

        try {
            EmbeddingR embeddingR = langEmbeddingService.embeddingText(
                    new ChatReq().setMessage(data.getContent())
                            .setDocsName(data.getType())
                            .setDocsId(data.getId())
                            .setKnowledgeId(data.getKnowledgeId()));

            aigcKnowledgeService.addDocsSlice(new AigcDocsSlice()
                    .setKnowledgeId(data.getKnowledgeId())
                    .setDocsId(data.getId())
                    .setVectorId(embeddingR.getVectorId())
                    .setName(data.getName())
                    .setContent(embeddingR.getText())
            );

            aigcKnowledgeService.updateDocs(new AigcDocs().setId(data.getId()).setSliceStatus(true).setSliceNum(1));
        } catch (Exception e) {
            e.printStackTrace();

            // del data
            aigcKnowledgeService.removeSlicesOfDoc(data.getId());
        }
        return CommonResult.success("success");
    }

    @PostMapping("/docs/{knowledgeId}")
    @PreAuthorize("@ss.hasPermission('aigc:embedding:docs')")
    @TenantIgnore
    public CommonResult docs(MultipartFile file, @PathVariable String knowledgeId) {
        // 1.1 获得用户信息
        String userId = getLoginUserId().toString();
        AigcOss oss = aigcOssService.upload(file, userId);
        AigcDocs data = new AigcDocs()
                .setName(oss.getOriginalFilename())
                .setSliceStatus(false)
                .setUrl(oss.getUrl())
                .setSize(file.getSize())
                .setType(EmbedConst.ORIGIN_TYPE_UPLOAD)
                .setKnowledgeId(knowledgeId);
        aigcKnowledgeService.addDocs(data);
        TaskManager.submitTask(userId, Executors.callable(() -> {
            embeddingService.embedDocsSlice(data, oss.getUrl());
        }));
        return CommonResult.success("success");
    }

    @GetMapping("/re-embed/{docsId}")
    @TenantIgnore
    public CommonResult reEmbed(@PathVariable String docsId) {
        String userId = getLoginUserId().toString();
        AigcDocs docs = aigcDocsMapper.selectById(docsId);
        if (docs == null) {
            throw new ServiceException(-1003,"没有查询到文档数据");
        }
        if (EmbedConst.ORIGIN_TYPE_INPUT.equals(docs.getType())) {
            text(docs);
        }
        if (EmbedConst.ORIGIN_TYPE_UPLOAD.equals(docs.getType())) {
            // clear before re-embed
            embeddingService.clearDocSlices(docsId);
            TaskManager.submitTask(userId, Executors.callable(() -> {
                embeddingService.embedDocsSlice(docs, docs.getUrl());
            }));
        }
        return CommonResult.success("success");
    }

    @PostMapping("/search")
    @TenantIgnore
    public CommonResult search(@RequestBody AigcDocs data) {
        return CommonResult.success(embeddingService.search(data));
    }
}
