package tech.easyflow.admin.controller.ai;

import cn.dev33.satoken.annotation.SaCheckPermission;
import com.easyagents.core.document.Document;
import com.easyagents.rag.retrieval.RagRetrievalMetadataKeys;
import com.mybatisflex.core.paginate.Page;
import com.mybatisflex.core.query.QueryWrapper;
import org.springframework.util.StringUtils;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
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.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import tech.easyflow.admin.controller.ai.support.AiResourceCreatorNameSupport;
import tech.easyflow.ai.permission.KnowledgeVisibilityQueryHelper;
import tech.easyflow.ai.documentimport.DocumentImportDtos;
import tech.easyflow.ai.dto.KnowledgeSearchResultItem;
import tech.easyflow.ai.entity.BotDocumentCollection;
import tech.easyflow.ai.entity.DocumentCollection;
import tech.easyflow.ai.entity.Model;
import tech.easyflow.ai.enums.PublishStatus;
import tech.easyflow.ai.publish.KnowledgePublishAppService;
import tech.easyflow.ai.service.AiResourceApprovalStateService;
import tech.easyflow.ai.vo.OfflineImpactCheckVo;
import tech.easyflow.approval.entity.vo.ApprovalActionResult;
import tech.easyflow.ai.rag.KnowledgeRetrievalRequest;
import tech.easyflow.ai.rag.KnowledgeRetrievalModes;
import tech.easyflow.ai.service.BotDocumentCollectionService;
import tech.easyflow.ai.service.DocumentChunkService;
import tech.easyflow.ai.service.DocumentCollectionService;
import tech.easyflow.ai.service.ModelService;
import tech.easyflow.common.domain.Result;
import tech.easyflow.common.web.controller.BaseCurdController;
import tech.easyflow.common.web.exceptions.BusinessException;
import tech.easyflow.common.web.jsonbody.JsonBody;
import tech.easyflow.system.enums.CategoryResourceType;
import tech.easyflow.system.enums.ResourceAction;
import tech.easyflow.system.enums.ResourceLookup;
import tech.easyflow.system.enums.VisibilityScope;
import tech.easyflow.system.permission.resource.RequireResourceAccess;
import tech.easyflow.system.service.ResourceAccessService;

import javax.annotation.Resource;
import java.io.Serializable;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.math.RoundingMode;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;

/**
 * 控制层。
 *
 * @author michael
 * @since 2024-08-23
 */
@RestController
@RequestMapping("/api/v1/documentCollection")
public class DocumentCollectionController extends BaseCurdController<DocumentCollectionService, DocumentCollection> {

    private final DocumentChunkService chunkService;
    private final ModelService llmService;

    @Resource
    private BotDocumentCollectionService botDocumentCollectionService;
    @Resource
    private ResourceAccessService resourceAccessService;
    @Resource
    private KnowledgeVisibilityQueryHelper knowledgeVisibilityQueryHelper;
    @Resource
    private KnowledgePublishAppService knowledgePublishAppService;
    @Resource
    private AiResourceApprovalStateService aiResourceApprovalStateService;
    @Resource
    private AiResourceCreatorNameSupport aiResourceCreatorNameSupport;

    public DocumentCollectionController(DocumentCollectionService service, DocumentChunkService chunkService, ModelService llmService) {
        super(service);
        this.chunkService = chunkService;
        this.llmService = llmService;
    }

    @Override
    protected Result<?> onSaveOrUpdateBefore(DocumentCollection entity, boolean isSave) {
        normalizeVisibilityScope(entity, isSave);
        if (!isSave && entity.getId() != null) {
            DocumentCollection existed = requireKnowledge(String.valueOf(entity.getId()));
            resourceAccessService.assertAccess(CategoryResourceType.KNOWLEDGE, existed, ResourceAction.MANAGE, "无权限管理知识库");
        }

        String alias = entity.getAlias();
        String collectionType = entity.getCollectionType();
        if (isSave && !StringUtils.hasLength(collectionType)) {
            entity.setCollectionType(DocumentCollection.TYPE_DOCUMENT);
        } else if (!DocumentCollection.TYPE_DOCUMENT.equalsIgnoreCase(collectionType)
                && !DocumentCollection.TYPE_FAQ.equalsIgnoreCase(collectionType)) {
            if (StringUtils.hasLength(collectionType)) {
                throw new BusinessException("知识库类型不正确");
            }
        } else {
            entity.setCollectionType(collectionType.toUpperCase());
        }

        if (StringUtils.hasLength(alias)){
            DocumentCollection knowledge = service.getByAlias(alias);

            if (knowledge != null && isSave){
                throw new BusinessException("别名已存在！");
            }

            if (knowledge != null && knowledge.getId().compareTo(entity.getId()) != 0){
                throw new BusinessException("别名已存在！");
            }

        } else {
            entity.setAlias(null);
        }


        if (isSave){
            Map<String, Object> options = entity.getOptions() == null
                ? new HashMap<>()
                : new HashMap<>(entity.getOptions());
            options.putIfAbsent(DocumentCollection.KEY_CAN_UPDATE_EMBEDDING_MODEL, true);
            options.putIfAbsent(DocumentCollection.KEY_RERANK_ENABLE, entity.getRerankModelId() != null);
            entity.setOptions(options);
        }
        normalizeInfrastructureFields(entity, isSave);
        return super.onSaveOrUpdateBefore(entity, isSave);
    }

    @GetMapping("search")
    @SaCheckPermission("/api/v1/documentCollection/query")
    @RequireResourceAccess(
            resource = CategoryResourceType.KNOWLEDGE,
            action = ResourceAction.READ,
            lookup = ResourceLookup.KNOWLEDGE_ID,
            idExpr = "#knowledgeId",
            denyMessage = "无权限访问知识库"
    )
    public Result<List<KnowledgeSearchResultItem>> search(@RequestParam BigInteger knowledgeId,
                                                          @RequestParam String keyword,
                                                          @RequestParam(required = false) String retrievalMode,
                                                          @RequestParam(required = false) Integer docRecallMaxNum,
                                                          @RequestParam(required = false) Double simThreshold) {
        KnowledgeRetrievalRequest request = new KnowledgeRetrievalRequest();
        request.setKnowledgeId(knowledgeId);
        request.setQuery(keyword);
        request.setLimit(docRecallMaxNum);
        request.setMinSimilarity(simThreshold);
        request.setRetrievalMode(KnowledgeRetrievalModes.parse(retrievalMode));
        request.setCallerType("API");
        request.setCallerId(String.valueOf(knowledgeId));
        return Result.ok(toKnowledgeSearchResult(service.search(request)));
    }


    @Override
    protected Result<Void> onRemoveBefore(Collection<Serializable> ids) {
        for (Serializable id : ids) {
            DocumentCollection collection = requireKnowledge(String.valueOf(id));
            resourceAccessService.assertAccess(CategoryResourceType.KNOWLEDGE, collection, ResourceAction.MANAGE, "无权限管理知识库");
        }

        QueryWrapper queryWrapper = QueryWrapper.create();
        queryWrapper.in(BotDocumentCollection::getDocumentCollectionId, ids);

        boolean exists = botDocumentCollectionService.exists(queryWrapper);
        if (exists){
            throw new BusinessException("此知识库还关联着bot，请先取消关联！");
        }

        return null;
    }

    @Override
    @RequireResourceAccess(
            resource = CategoryResourceType.KNOWLEDGE,
            action = ResourceAction.READ,
            lookup = ResourceLookup.KNOWLEDGE_ID_OR_SLUG,
            idExpr = "#id",
            denyMessage = "无权限访问知识库"
    )
    public Result<DocumentCollection> detail(String id) {
        DocumentCollection detail = service.getDetail(id);
        aiResourceApprovalStateService.fillKnowledgeApprovalState(detail);
        return Result.ok(detail);
    }

    @GetMapping("modelList")
    @SaCheckPermission("/api/v1/documentCollection/query")
    public Result<List<Model>> modelList(Model entity, Boolean asTree, String sortKey, String sortType) {
        return Result.ok(llmService.listSelectableModels(entity, asTree, sortKey, sortType));
    }

    /**
     * 提交发布审批。
     *
     * @param id 知识库 ID
     * @return 审批实例 ID
     */
    @PostMapping("/submitPublishApproval")
    @SaCheckPermission("/api/v1/documentCollection/save")
    public Result<BigInteger> submitPublishApproval(@JsonBody("id") BigInteger id) {
        return buildApprovalActionResult(
                knowledgePublishAppService.submitPublishApproval(id),
                "已提交发布审批",
                "已直接发布"
        );
    }

    /**
     * 提交下线审批。
     *
     * @param id 知识库 ID
     * @return 审批实例 ID
     */
    @PostMapping("/submitOfflineApproval")
    @SaCheckPermission("/api/v1/documentCollection/save")
    public Result<BigInteger> submitOfflineApproval(@JsonBody("id") BigInteger id) {
        return buildApprovalActionResult(
                knowledgePublishAppService.submitOfflineApproval(id),
                "已提交下线审批",
                "已直接下线"
        );
    }

    /**
     * 检查知识库下线影响。
     *
     * @param id 知识库 ID
     * @return 下线影响结果
     */
    @GetMapping("/offlineImpactCheck")
    @SaCheckPermission("/api/v1/documentCollection/save")
    @RequireResourceAccess(
            resource = CategoryResourceType.KNOWLEDGE,
            action = ResourceAction.MANAGE,
            lookup = ResourceLookup.KNOWLEDGE_ID,
            idExpr = "#id",
            denyMessage = "无权限管理知识库"
    )
    public Result<OfflineImpactCheckVo> offlineImpactCheck(@RequestParam BigInteger id) {
        return Result.ok(knowledgePublishAppService.checkOfflineImpact(id));
    }

    /**
     * 提交删除审批。
     *
     * @param id 知识库 ID
     * @return 审批实例 ID
     */
    @PostMapping("/submitDeleteApproval")
    @SaCheckPermission("/api/v1/documentCollection/remove")
    public Result<BigInteger> submitDeleteApproval(@JsonBody("id") BigInteger id) {
        return buildApprovalActionResult(
                knowledgePublishAppService.submitDeleteApproval(id),
                "已提交删除审批",
                "已直接删除"
        );
    }

    @PostMapping("splitterProfile/save")
    @SaCheckPermission("/api/v1/documentCollection/save")
    @RequireResourceAccess(
            resource = CategoryResourceType.KNOWLEDGE,
            action = ResourceAction.MANAGE,
            lookup = ResourceLookup.KNOWLEDGE_ID,
            idExpr = "#request.knowledgeId",
            denyMessage = "无权限管理知识库"
    )
    public Result<Boolean> saveSplitterProfile(@JsonBody DocumentImportDtos.SplitterProfileSaveRequest request) {
        if (request.getKnowledgeId() == null) {
            throw new BusinessException("知识库ID不能为空");
        }
        DocumentCollection collection = service.getById(request.getKnowledgeId());
        if (collection == null) {
            throw new BusinessException("知识库不存在");
        }
        if (collection.isFaqCollection()) {
            throw new BusinessException("FAQ知识库不支持文档导入策略");
        }
        Map<String, Object> options = collection.getOptions() == null
                ? new HashMap<>()
                : new HashMap<>(collection.getOptions());
        options.put(DocumentCollection.KEY_SPLITTER_DEFAULT_STRATEGY, request.getDefaultStrategyCode());
        options.put(DocumentCollection.KEY_SPLITTER_AUTO_RECOMMEND_ENABLED, request.getAutoRecommendEnabled());
        options.put(DocumentCollection.KEY_SPLITTER_FALLBACK_STRATEGY, request.getFallbackStrategyCode());
        options.put(DocumentCollection.KEY_SPLITTER_STRATEGY_PROFILES, request.getStrategyProfiles());

        DocumentCollection update = new DocumentCollection();
        update.setId(collection.getId());
        update.setOptions(options);
        return Result.ok(service.updateById(update));
    }

    @Override
    public Result<List<DocumentCollection>> list(DocumentCollection entity, Boolean asTree, String sortKey, String sortType) {
        QueryWrapper queryWrapper = QueryWrapper.create(entity, buildOperators(entity));
        knowledgeVisibilityQueryHelper.applyReadableAccess(queryWrapper);
        applyPublishedOnlyFilter(queryWrapper);
        queryWrapper.orderBy(buildOrderBy(sortKey, sortType, getDefaultOrderBy()));
        List<DocumentCollection> collections = service.list(queryWrapper);
        aiResourceApprovalStateService.fillKnowledgeApprovalState(collections);
        return Result.ok(collections);
    }

    @Override
    protected Page<DocumentCollection> queryPage(Page<DocumentCollection> page, QueryWrapper queryWrapper) {
        knowledgeVisibilityQueryHelper.applyReadableAccess(queryWrapper);
        applyPublishedOnlyFilter(queryWrapper);
        Page<DocumentCollection> result = super.queryPage(page, queryWrapper);
        aiResourceApprovalStateService.fillKnowledgeApprovalState(result.getRecords());
        aiResourceCreatorNameSupport.fillDocumentCollectionCreatorNames(result.getRecords());
        return result;
    }

    @Override
    @PostMapping("remove")
    public Result<?> remove(@JsonBody(value = "id", required = true) Serializable id) {
        return Result.fail(1, "请使用发布状态操作删除知识库");
    }

    private Result<BigInteger> buildApprovalActionResult(ApprovalActionResult actionResult,
                                                         String approvalMessage,
                                                         String directMessage) {
        return Result.ok(
                actionResult.isApprovalRequired() ? approvalMessage : directMessage,
                actionResult.getInstanceId()
        );
    }

    private void normalizeVisibilityScope(DocumentCollection entity, boolean isSave) {
        if (entity == null) {
            return;
        }
        if (!StringUtils.hasLength(entity.getVisibilityScope())) {
            if (isSave) {
                entity.setVisibilityScope(VisibilityScope.PRIVATE.name());
            }
            return;
        }
        entity.setVisibilityScope(VisibilityScope.from(entity.getVisibilityScope()).name());
    }

    private DocumentCollection requireKnowledge(String idOrAlias) {
        DocumentCollection collection = service.getDetail(idOrAlias);
        if (collection == null) {
            throw new BusinessException("知识库不存在");
        }
        return collection;
    }

    private void normalizeInfrastructureFields(DocumentCollection entity, boolean isSave) {
        if (entity == null) {
            return;
        }
        if (isSave) {
            entity.setVectorStoreEnable(true);
            entity.setVectorStoreType("milvus");
            entity.setSearchEngineEnable(true);
            entity.setVectorStoreCollection(generateVectorCollectionName());
            return;
        }
        if (entity.getVectorStoreEnable() != null) {
            entity.setVectorStoreEnable(true);
        }
        if (entity.getVectorStoreType() != null) {
            entity.setVectorStoreType("milvus");
        }
        if (entity.getSearchEngineEnable() != null) {
            entity.setSearchEngineEnable(true);
        }
    }

    private String generateVectorCollectionName() {
        return "kb_" + UUID.randomUUID().toString().replace("-", "").substring(0, 28);
    }

    private void applyPublishedOnlyFilter(QueryWrapper queryWrapper) {
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (attributes == null) {
            return;
        }
        String publishedOnly = attributes.getRequest().getParameter("publishedOnly");
        if ("true".equalsIgnoreCase(publishedOnly)) {
            queryWrapper.eq("publish_status", PublishStatus.PUBLISHED.getCode());
        }
    }

    private List<KnowledgeSearchResultItem> toKnowledgeSearchResult(List<Document> documents) {
        List<KnowledgeSearchResultItem> results = new ArrayList<>();
        if (documents == null) {
            return results;
        }
        for (int index = 0; index < documents.size(); index++) {
            Document document = documents.get(index);
            if (document == null) {
                continue;
            }
            KnowledgeSearchResultItem item = new KnowledgeSearchResultItem();
            item.setSorting(index + 1);
            item.setContent(document.getContent());
            item.setRenderMarkdown(readMetadataAsString(document, "renderMarkdown"));
            item.setSourceFileName(readMetadataAsString(document, "sourceFileName"));
            item.setScore(roundScore(document.getScore()));
            item.setHitSource(readMetadataAsString(document, RagRetrievalMetadataKeys.HIT_SOURCE));
            item.setVectorScore(roundScore(readMetadataAsDouble(document, RagRetrievalMetadataKeys.VECTOR_SCORE)));
            item.setKeywordScore(roundScore(readMetadataAsDouble(document, RagRetrievalMetadataKeys.KEYWORD_SCORE)));
            results.add(item);
        }
        return results;
    }

    private String readMetadataAsString(Document document, String key) {
        Object value = document == null ? null : document.getMetadata(key);
        if (value == null) {
            return null;
        }
        String text = String.valueOf(value);
        return StringUtils.hasText(text) ? text : null;
    }

    private Double readMetadataAsDouble(Document document, String key) {
        Object value = document == null ? null : document.getMetadata(key);
        if (value instanceof Number) {
            return ((Number) value).doubleValue();
        }
        if (value instanceof String && StringUtils.hasText((String) value)) {
            try {
                return Double.valueOf((String) value);
            } catch (NumberFormatException ignore) {
                return null;
            }
        }
        return null;
    }

    private Double roundScore(Double value) {
        if (value == null) {
            return null;
        }
        return new BigDecimal(String.valueOf(value))
            .setScale(4, RoundingMode.HALF_UP)
            .doubleValue();
    }
}
