RAG/src/main/java/com/wok/supportbot/controller/AiController.java

package com.wok.supportbot.controller;

import com.wok.supportbot.app.AssistantApp;
import com.wok.supportbot.config.RoleAccessConfig;
import com.wok.supportbot.service.ConversationService;
import com.wok.supportbot.service.CustomerServiceRoleService;
import com.wok.supportbot.service.CustomerServiceRoleService.RoleScope;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.DeleteMapping;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;


@RestController
@Slf4j
@RequestMapping("/ai")
public class AiController {

    @Resource
    private AssistantApp assistantApp;
    @Resource
    private CustomerServiceRoleService customerServiceRoleService;
    @Resource
    private ConversationService conversationService;
    @Resource
    private RoleAccessConfig roleAccessConfig;


    /**
     * 同步调用 AI 智能客服应用
     *
     * @param message      用户消息
     * @param chatId       会话ID
     * @param roleId       客服角色ID（可选）；命中角色时由服务端强制套用该角色人设
     * @param systemPrompt 角色人设兜底（仅在未命中角色时生效）
     * @return AI 回答
     */
    @GetMapping("/assistant_app/chat/sync")
    public String doChatWithAssistantAppSync(String message, String chatId, Long roleId, String accountId, String systemPrompt) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        bindConversation(chatId, context);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        String result = assistantApp.doChat(message, chatId, resolveSystemPrompt(scope, systemPrompt), scope.allowedMcpTools());
        return result;
    }

    /**
     * SSE 流式调用 AI 智能客服应用
     * 返回Flux 响应式؜对象，并且添加 SSE 对应的 MediaType
     *
     * @param message
     * @param chatId
     * @return
     */
    @GetMapping(value = "/assistant_app/chat/sse", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> doChatWithLoveAppSSE(String message, String chatId, Long roleId, String accountId, String systemPrompt) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        bindConversation(chatId, context);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        return assistantApp.doChatByStream(message, chatId, resolveSystemPrompt(scope, systemPrompt), scope.allowedMcpTools());
    }

    /**
     * SSE 流式调用 AI 智能客服应用
     * 返回 Flux 对象，并且؜设置泛型为 ServerSentEvent。使用这种方式可以省略 MediaType
     *
     * @param message
     * @param chatId
     * @return
     */
    @GetMapping(value = "/assistant_app/chat/server_sent_event")
    public Flux<ServerSentEvent<String>> doChatWithAssistantAppServerSentEvent(String message, String chatId, Long roleId, String accountId, String systemPrompt) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        bindConversation(chatId, context);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        return assistantApp.doChatByStream(message, chatId, resolveSystemPrompt(scope, systemPrompt), scope.allowedMcpTools())
                .map(chunk -> ServerSentEvent.<String>builder()
                        .data(chunk)
                        .build());
    }

    /**
     * SSE 流式调用 AI 智能客服应用
     * 使用 SSEEmiter，؜通过 send 方法持续向 SseEmitter 发送消息
     *
     * @param message
     * @param chatId
     * @return
     */
    @GetMapping(value = "/assistant_app/chat/sse_emitter")
    public SseEmitter doChatWithAssistantAppServerSseEmitter(String message, String chatId, Long roleId, String accountId, String systemPrompt) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        bindConversation(chatId, context);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        // 创建一个超时时间较长的 SseEmitter
        SseEmitter sseEmitter = new SseEmitter(180000L); // 3 分钟超时
        // 获取 Flux 响应式数据流并且直接通过订阅推送给 SseEmitter
        assistantApp.doChatByStream(message, chatId, resolveSystemPrompt(scope, systemPrompt), scope.allowedMcpTools())
                .subscribe(chunk -> {
                    try {
                        sseEmitter.send(chunk);
                    } catch (IOException e) {
                        sseEmitter.completeWithError(e);
                    }
                }, sseEmitter::completeWithError, sseEmitter::complete);
        // 返回
        return sseEmitter;
    }

    /**
     * RAG 知识库同步对话（支持查询重写策略）
     *
     * @param message         用户消息
     * @param chatId          会话ID
     * @param rewriteStrategy 查询重写策略（可选）：NONE/REWRITE/TRANSLATION/COMPRESSION/MULTI_QUERY，默认为 MULTI_QUERY
     * @param roleId          客服角色ID（可选）；命中角色时由服务端强制限定检索范围与人设，客户端无法跨域
     * @param categoryId      单个分类过滤（仅未命中角色时生效）
     * @param categoryIds     多个分类过滤，逗号分隔（仅未命中角色时生效）
     * @param systemPrompt    角色人设兜底（仅未命中角色时生效）
     * @return AI 回答
     */
    @GetMapping("/assistant_app/chat/rag/sync")
    public String doChatWithRagSync(String message, String chatId, String rewriteStrategy, Long roleId, String accountId, Long categoryId, String categoryIds, String systemPrompt) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        bindConversation(chatId, context);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        String sys = resolveSystemPrompt(scope, systemPrompt);
        // 严格隔离：未授权任何知识库的角色退回普通对话，绝不检索 KB
        if (isKbDenied(scope) || shouldBypassKnowledgeRetrieval(message)) {
            return assistantApp.doChat(message, chatId, sys, scope.allowedMcpTools());
        }
        try {
            return assistantApp.doChatWithRagStrategy(
                    message, chatId, normalizeStrategy(rewriteStrategy),
                    resolveCategoryIds(scope, categoryId, categoryIds), sys, scope.allowedMcpTools());
        } catch (Exception e) {
            log.error("RAG 对话失败 [strategy={}, chatId={}]: {}", rewriteStrategy, chatId, e.getMessage(), e);
            return "抱歉，知识库检索出现异常，请稍后重试。";
        }
    }

    /**
     * RAG 知识库流式对话（SSE，支持查询重写策略）
     *
     * @param message         用户消息
     * @param chatId          会话ID
     * @param rewriteStrategy 查询重写策略（可选）：NONE/REWRITE/TRANSLATION/COMPRESSION/MULTI_QUERY，默认为 MULTI_QUERY
     * @param roleId          客服角色ID（可选）；命中角色时由服务端强制限定检索范围与人设，客户端无法跨域
     * @param categoryId      单个分类过滤（仅未命中角色时生效）
     * @param categoryIds     多个分类过滤，逗号分隔（仅未命中角色时生效）
     * @param systemPrompt    角色人设兜底（仅未命中角色时生效）
     * @return 流式 AI 回答
     */
    @GetMapping(value = "/assistant_app/chat/rag/sse", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> doChatWithRagSSE(String message, String chatId, String rewriteStrategy, Long roleId, String accountId, Long categoryId, String categoryIds, String systemPrompt) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        bindConversation(chatId, context);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        String sys = resolveSystemPrompt(scope, systemPrompt);
        // 严格隔离：未授权任何知识库的角色退回普通流式对话，绝不检索 KB
        if (isKbDenied(scope) || shouldBypassKnowledgeRetrieval(message)) {
            return assistantApp.doChatByStream(message, chatId, sys, scope.allowedMcpTools());
        }
        return assistantApp.doChatWithRagStrategyByStream(
                message, chatId, normalizeStrategy(rewriteStrategy),
                resolveCategoryIds(scope, categoryId, categoryIds), sys, scope.allowedMcpTools());
    }

    /**
     * RAG 引用来源：返回本次问题命中的知识库片段（不生成回答），供回答下方展示来源。
     * 复用与 RAG 回答相同的角色范围、分类过滤与查询改写策略，确保来源即答案所依据的片段。
     */
    @GetMapping("/assistant_app/rag/sources")
    public Map<String, Object> getRagSources(String message, String chatId, String rewriteStrategy,
                                             Long roleId, String accountId, Long categoryId, String categoryIds) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        if (message == null || message.isBlank() || isKbDenied(scope) || shouldBypassKnowledgeRetrieval(message)) {
            return Map.of("success", true, "data", List.of());
        }
        try {
            List<Long> cats = resolveCategoryIds(scope, categoryId, categoryIds);
            List<Document> docs = assistantApp.retrieveRagSources(message, chatId, normalizeStrategy(rewriteStrategy), cats);
            List<Map<String, Object>> out = new ArrayList<>();
            for (Document doc : docs) {
                Map<String, Object> meta = doc.getMetadata();
                Map<String, Object> item = new LinkedHashMap<>();
                item.put("documentId", meta.get("documentId"));
                item.put("title", meta.get("title"));
                item.put("sourceName", meta.get("sourceName"));
                item.put("chunkIndex", meta.get("chunkIndex"));
                item.put("score", meta.get("distance"));
                String text = doc.getText();
                item.put("snippet", text != null && text.length() > 160 ? text.substring(0, 160) + "…" : text);
                out.add(item);
            }
            return Map.of("success", true, "data", out);
        } catch (Exception e) {
            log.error("获取 RAG 引用来源失败 [strategy={}]: {}", rewriteStrategy, e.getMessage(), e);
            return Map.of("success", true, "data", List.of());
        }
    }

    // ==================== SDK 会话管理接口（带账户归属校验） ====================

    /**
     * SDK 获取会话列表（分页，强制按账户 + 角色过滤）
     * 仅返回属于指定 accountId + roleId 的会话，防止跨账户数据泄露。
     *
     * @param page      页码（默认1）
     * @param size      每页大小（默认20）
     * @param accountId 外部账户 ID（SDK 的 userId，必传）
     * @param roleId    角色 ID（SDK 的 integrateId，必传）
     * @return 分页会话列表
     */
    @GetMapping("/sdk/conversation/list")
    public Map<String, Object> sdkListConversations(
            @RequestParam(defaultValue = "1") int page,
            @RequestParam(defaultValue = "20") int size,
            @RequestParam String accountId,
            @RequestParam Long roleId) {
        try {
            Map<String, Object> result = conversationService.listConversationsForSdk(page, size, accountId, roleId);
            Map<String, Object> data = new java.util.HashMap<>();
            data.put("success", true);
            data.put("data", result.get("records"));
            data.put("total", result.get("total"));
            data.put("page", result.get("page"));
            data.put("size", result.get("size"));
            data.put("pages", result.get("pages"));
            return data;
        } catch (Exception e) {
            log.error("SDK 会话列表查询失败: accountId={}, roleId={}", accountId, roleId, e);
            return Map.of("success", false, "message", "查询失败：" + e.getMessage());
        }
    }

    /**
     * SDK 获取会话消息（含所有权校验）
     * 先验证会话是否属于指定 accountId + roleId，通过后才返回消息。
     *
     * @param conversationId 会话 ID
     * @param accountId      外部账户 ID（必传）
     * @param roleId         角色 ID（必传）
     * @return 消息列表
     */
    @GetMapping("/sdk/conversation/{id}/messages")
    public Map<String, Object> sdkGetConversationMessages(
            @PathVariable("id") String conversationId,
            @RequestParam String accountId,
            @RequestParam Long roleId) {
        try {
            List<Map<String, Object>> messages =
                    conversationService.getConversationMessagesForSdk(conversationId, accountId, roleId);
            return Map.of("success", true, "data", messages, "total", messages.size());
        } catch (IllegalArgumentException e) {
            return Map.of("success", false, "message", e.getMessage());
        } catch (Exception e) {
            log.error("SDK 会话消息查询失败: conversationId={}, accountId={}", conversationId, accountId, e);
            return Map.of("success", false, "message", "查询失败：" + e.getMessage());
        }
    }

    /**
     * SDK 删除会话（含所有权校验）
     * 仅允许删除属于自己的会话。
     *
     * @param conversationId 会话 ID
     * @param accountId      外部账户 ID（必传）
     * @param roleId         角色 ID（必传）
     * @return 删除结果
     */
    @DeleteMapping("/sdk/conversation/{id}")
    public Map<String, Object> sdkDeleteConversation(
            @PathVariable("id") String conversationId,
            @RequestParam String accountId,
            @RequestParam Long roleId) {
        try {
            int count = conversationService.deleteConversationForSdk(conversationId, accountId, roleId);
            return Map.of("success", true, "message", "会话删除成功", "deletedMessages", count);
        } catch (IllegalArgumentException e) {
            return Map.of("success", false, "message", e.getMessage());
        } catch (Exception e) {
            log.error("SDK 会话删除失败: conversationId={}, accountId={}", conversationId, accountId, e);
            return Map.of("success", false, "message", "删除失败：" + e.getMessage());
        }
    }

    /**
     * SDK 导出会话（含所有权校验，返回文本内容）
     * 通过 accountId + roleId 验证会话归属后导出。
     *
     * @param conversationId 会话 ID
     * @param accountId      外部账户 ID（必传）
     * @param roleId         角色 ID（必传）
     * @return 导出的会话文本
     */
    @GetMapping(value = "/sdk/conversation/{id}/export", produces = MediaType.TEXT_PLAIN_VALUE + ";charset=UTF-8")
    public String sdkExportConversation(
            @PathVariable("id") String conversationId,
            @RequestParam String accountId,
            @RequestParam Long roleId) {
        try {
            conversationService.getConversationMessagesForSdk(conversationId, accountId, roleId);
            return conversationService.exportConversation(conversationId);
        } catch (IllegalArgumentException e) {
            return "导出失败：" + e.getMessage();
        } catch (Exception e) {
            log.error("SDK 会话导出失败: conversationId={}, accountId={}", conversationId, accountId, e);
            return "导出失败：" + e.getMessage();
        }
    }

    /**
     * SDK 截断会话（含所有权校验）
     * 用于"编辑历史消息并重发"场景，仅允许截断属于自己的会话。
     *
     * @param conversationId 会话 ID
     * @param body           { "userTurn": 第几条用户消息, "accountId": "xxx", "roleId": 1 }
     * @return 逻辑删除的消息条数
     */
    @PostMapping("/sdk/conversation/{id}/truncate")
    public Map<String, Object> sdkTruncateConversation(
            @PathVariable("id") String conversationId,
            @RequestBody java.util.Map<String, Object> body) {
        try {
            Object rawTurn = body.get("userTurn");
            int userTurn = rawTurn instanceof Number n ? n.intValue() : Integer.parseInt(String.valueOf(rawTurn));
            String accountId = (String) body.get("accountId");
            Long roleId = body.get("roleId") instanceof Number n ? n.longValue() : null;
            int deleted = conversationService.truncateFromUserTurnForSdk(
                    conversationId, userTurn, accountId, roleId);
            return Map.of("success", true, "deletedMessages", deleted);
        } catch (IllegalArgumentException e) {
            return Map.of("success", false, "message", e.getMessage());
        } catch (Exception e) {
            log.error("SDK 会话截断失败: conversationId={}", conversationId, e);
            return Map.of("success", false, "message", "截断失败：" + e.getMessage());
        }
    }

    /**
     * 严格隔离模式下，命中角色但未绑定任何知识库分类 ⇒ 拒绝检索任何 KB。
     * 非严格模式下永远返回 false（未绑定 = 可检索全部，沿用通用角色语义）。
     */
    private boolean isKbDenied(RoleScope scope) {
        return roleAccessConfig.isStrictIsolation() && scope.hasRole() && scope.categoryIds().isEmpty();
    }

    /**
     * 寒暄、感谢、告别这类短消息没有知识库检索意图，直接走普通对话。
     * 否则向量库会被迫召回一个“最相似但实际无关”的片段，导致回答下方出现误导性来源。
     */
    private boolean shouldBypassKnowledgeRetrieval(String message) {
        if (!StringUtils.hasText(message)) {
            return true;
        }
        String normalized = message.trim()
                .toLowerCase(Locale.ROOT)
                .replaceAll("[\\s，。！？!?,.;；：:、~～…]+", "");
        if (normalized.isEmpty()) {
            return true;
        }
        if (normalized.length() > 12) {
            return false;
        }
        return List.of(
                "你好", "您好", "hello", "hi", "哈喽", "嗨", "在吗",
                "早上好", "上午好", "中午好", "下午好", "晚上好",
                "谢谢", "感谢", "多谢", "好的", "好", "嗯", "嗯嗯",
                "再见", "拜拜", "辛苦了"
        ).contains(normalized);
    }

    /** 未指定策略时默认 MULTI_QUERY（多路扩展）。 */
    private String normalizeStrategy(String rewriteStrategy) {
        return (rewriteStrategy != null && !rewriteStrategy.isEmpty()) ? rewriteStrategy : "MULTI_QUERY";
    }

    /**
     * 命中角色时用 roleId 套用角色人设和知识库范围。
     * accountId 是 SDK 宿主侧的外部用户标识，后端不再创建或查询本地账号，也不允许它覆盖 roleId。
     */
    private AccountRoleContext resolveAccountRole(String accountId, Long fallbackRoleId) {
        String externalAccountId = StringUtils.hasText(accountId) ? accountId.trim() : null;
        return new AccountRoleContext(externalAccountId, fallbackRoleId);
    }
    private void bindConversation(String chatId, AccountRoleContext context) {
        conversationService.bindConversation(chatId, context.accountId(), context.roleId());
    }

    private String resolveSystemPrompt(RoleScope scope, String fallbackSystemPrompt) {
        if (scope.hasRole() && StringUtils.hasText(scope.systemPrompt())) {
            return scope.systemPrompt();
        }
        return fallbackSystemPrompt;
    }

    /**
     * 解析最终的检索分类范围。
     * 命中角色时，强制使用角色绑定的分类（忽略客户端传入），防止越权跨域；
     * 未命中角色时，沿用客户端传入的分类（便于直接调试/无角色调用）。
     */
    private List<Long> resolveCategoryIds(RoleScope scope, Long categoryId, String categoryIds) {
        if (scope.hasRole()) {
            return scope.categoryIds();
        }
        List<Long> ids = parseCategoryIds(categoryIds);
        if (ids.isEmpty() && categoryId != null) {
            return List.of(categoryId);
        }
        return ids;
    }

    private List<Long> parseCategoryIds(String categoryIds) {
        if (categoryIds == null || categoryIds.isBlank()) {
            return Collections.emptyList();
        }
        return Arrays.stream(categoryIds.split(","))
                .map(String::trim)
                .filter(item -> !item.isEmpty())
                .map(Long::valueOf)
                .filter(id -> id > 0)
                .distinct()
                .toList();
    }

    private record AccountRoleContext(String accountId, Long roleId) {
    }

}