package tech.easyflow.admin.controller.ai;

import cn.dev33.satoken.annotation.SaCheckPermission;
import com.easyagents.core.model.embedding.EmbeddingModel;
import com.mybatisflex.core.paginate.Page;
import tech.easyflow.ai.entity.DocumentChunk;
import tech.easyflow.ai.entity.DocumentCollection;
import tech.easyflow.ai.entity.Model;
import tech.easyflow.ai.service.DocumentChunkService;
import tech.easyflow.ai.service.DocumentCollectionService;
import tech.easyflow.ai.service.ModelService;
import tech.easyflow.ai.support.DocumentStoreLifecycleSupport;
import tech.easyflow.common.annotation.UsePermission;
import tech.easyflow.common.domain.Result;
import tech.easyflow.common.web.controller.BaseCurdController;
import tech.easyflow.common.web.jsonbody.JsonBody;
import com.easyagents.core.document.Document;
import com.easyagents.core.store.DocumentStore;
import com.easyagents.core.store.StoreOptions;
import com.easyagents.core.store.StoreResult;
import jakarta.servlet.http.HttpServletRequest;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import tech.easyflow.system.enums.CategoryResourceType;
import tech.easyflow.system.enums.ResourceAction;
import tech.easyflow.system.enums.ResourceLookup;
import tech.easyflow.system.permission.resource.RequireResourceAccess;

import javax.annotation.Resource;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 *  控制层。
 *
 * @author michael
 * @since 2024-08-23
 */
@RestController
@RequestMapping("/api/v1/documentChunk")
@UsePermission(moduleName = "/api/v1/documentCollection")
public class DocumentChunkController extends BaseCurdController<DocumentChunkService, DocumentChunk> {

    @Resource
    DocumentCollectionService documentCollectionService;

    @Resource
    ModelService modelService;

    @Resource
    DocumentChunkService documentChunkService;

    public DocumentChunkController(DocumentChunkService service) {
        super(service);
    }

    @GetMapping("page")
    @SaCheckPermission("/api/v1/documentCollection/query")
    @RequireResourceAccess(
            resource = CategoryResourceType.KNOWLEDGE,
            action = ResourceAction.READ,
            lookup = ResourceLookup.DOCUMENT_ID,
            idExpr = "#request.getParameter('documentId')",
            denyMessage = "无权限访问知识库"
    )
    @Override
    public Result<Page<DocumentChunk>> page(HttpServletRequest request, String sortKey, String sortType, Long pageNumber, Long pageSize) {
        return super.page(request, sortKey, sortType, pageNumber, pageSize);
    }

    @PostMapping("update")
    @SaCheckPermission("/api/v1/documentCollection/save")
    @RequireResourceAccess(
            resource = CategoryResourceType.KNOWLEDGE,
            action = ResourceAction.MANAGE,
            lookup = ResourceLookup.DOCUMENT_CHUNK_ID,
            idExpr = "#documentChunk.id",
            denyMessage = "无权限管理知识库"
    )
    public Result<?> update(@JsonBody DocumentChunk documentChunk) {
        boolean success = service.updateById(documentChunk);
        if (success){
            DocumentChunk record = documentChunkService.getById(documentChunk.getId());
            DocumentCollection knowledge = documentCollectionService.getById(record.getDocumentCollectionId());
            if (knowledge == null) {
                return Result.fail(1, "知识库不存在");
            }
            DocumentStore documentStore = knowledge.toDocumentStore();
            if (documentStore == null) {
                return Result.fail(2, "知识库没有配置向量库");
            }
            try {
                // 设置向量模型
                Model model = modelService.getModelInstance(knowledge.getVectorEmbedModelId());
                if (model == null) {
                    return Result.fail(3, "知识库没有配置向量模型");
                }
                EmbeddingModel embeddingModel = model.toEmbeddingModel();
                documentStore.setEmbeddingModel(embeddingModel);
                StoreOptions options = StoreOptions.ofCollectionName(knowledge.getVectorStoreCollection());
                Document document = Document.of(documentChunk.getContent());
                document.setId(documentChunk.getId());
                Map<String, Object> metadata = new HashMap<>();
                metadata.put("keywords", documentChunk.getMetadataKeyWords());
                metadata.put("questions", documentChunk.getMetadataQuestions());
                document.setMetadataMap(metadata);
                StoreResult result = documentStore.update(document, options); // 更新已有记录
                return Result.ok(result);
            } finally {
                DocumentStoreLifecycleSupport.closeQuietly(documentStore);
            }
        }
        return Result.ok(false);
    }

    @PostMapping("removeChunk")
    @SaCheckPermission("/api/v1/documentCollection/remove")
    @RequireResourceAccess(
            resource = CategoryResourceType.KNOWLEDGE,
            action = ResourceAction.MANAGE,
            lookup = ResourceLookup.DOCUMENT_CHUNK_ID,
            idExpr = "#chunkId",
            denyMessage = "无权限管理知识库"
    )
    public Result<?> remove(@JsonBody(value = "id", required = true) BigInteger chunkId) {
        DocumentChunk docChunk =  documentChunkService.getById(chunkId);
        if (docChunk == null) {
            return Result.fail(1, "记录不存在");
        }
        DocumentCollection knowledge = documentCollectionService.getById(docChunk.getDocumentCollectionId());
        if (knowledge == null) {
            return Result.fail(2, "知识库不存在");
        }
        DocumentStore documentStore = knowledge.toDocumentStore();
        if (documentStore == null) {
            return Result.fail(3, "知识库没有配置向量库");
        }
        try {
            // 设置向量模型
            Model model = modelService.getModelInstance(knowledge.getVectorEmbedModelId());
            if (model == null) {
                return Result.fail(4, "知识库没有配置向量模型");
            }
            EmbeddingModel embeddingModel = model.toEmbeddingModel();
            documentStore.setEmbeddingModel(embeddingModel);
            StoreOptions options = StoreOptions.ofCollectionName(knowledge.getVectorStoreCollection());
            List<BigInteger> deleteList = new ArrayList<>();
            deleteList.add(chunkId);
            documentStore.delete(deleteList, options);
            documentChunkService.removeChunk(knowledge, chunkId);

            return super.remove(chunkId);
        } finally {
            DocumentStoreLifecycleSupport.closeQuietly(documentStore);
        }
    }
}
