sundynix-agentix/sundynix-gateway/internal/handler/kb.go

package handler

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/gin-gonic/gin"

	"github.com/sundynix/sundynix-gateway/internal/blob"
	"github.com/sundynix/sundynix-gateway/internal/nats"
	"github.com/sundynix/sundynix-gateway/internal/store"
	"github.com/sundynix/sundynix-shared/contract"
)

// rawKB 规整知识库名（去空白，空则 default）—— 注册表里的展示名。
func rawKB(kb string) string {
	kb = strings.TrimSpace(kb)
	if kb == "" {
		return "default"
	}
	return kb
}

// scopedKB 把知识库名锁进当前用户作用域："owner/name"。
// owner 来自身份（X-User-ID），客户端只发库名、发不了 owner，故无法越权查到他人的库。
func scopedKB(c *gin.Context, kb string) string {
	return userID(c) + "/" + rawKB(kb)
}

// KbList: GET /api/v1/kb/list —— 当前用户的知识库列表（按 owner 隔离）。
func (h *Handler) KbList(c *gin.Context) {
	rows, err := h.db.ListKB(c.Request.Context(), userID(c))
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	out := make([]gin.H, 0, len(rows))
	for _, r := range rows {
		out = append(out, gin.H{"name": r.Name, "kind": r.Kind})
	}
	c.JSON(http.StatusOK, gin.H{"kbs": out})
}

// KbCreate: POST /api/v1/kb/create {name, kind} —— 新建知识库（folder/project/case/general）。
func (h *Handler) KbCreate(c *gin.Context) {
	var body struct {
		Name string `json:"name"`
		Kind string `json:"kind"`
	}
	if err := c.ShouldBindJSON(&body); err != nil || strings.TrimSpace(body.Name) == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "name required"})
		return
	}
	if err := h.db.EnsureKB(c.Request.Context(), userID(c), rawKB(body.Name), body.Kind); err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"name": rawKB(body.Name), "kind": body.Kind})
}

// KbIngest: POST /api/v1/kb/ingest —— 文本入库（异步，返回 job_id；进度经 SSE 看）。
func (h *Handler) KbIngest(c *gin.Context) {
	var body struct {
		KB   string `json:"kb"`
		Text string `json:"text"`
	}
	if err := c.ShouldBindJSON(&body); err != nil || body.Text == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "text required"})
		return
	}
	_ = h.db.EnsureKB(c.Request.Context(), userID(c), rawKB(body.KB), "general")
	job, err := h.enqueueIngest(c.Request.Context(), userID(c), rawKB(body.KB), scopedKB(c, body.KB), "", "", nil, body.Text)
	if err != nil {
		c.JSON(http.StatusServiceUnavailable, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusAccepted, gin.H{"job_id": job})
}

// KbSaveNote: POST /api/v1/kb/note {kb, name, content} —— 新建/编辑笔记。
// 立即落库（文库可见），并以 name 为 doc 重新入库（替换旧块，搜索/图谱同步）。返回 job_id。
func (h *Handler) KbSaveNote(c *gin.Context) {
	var body struct {
		KB      string `json:"kb"`
		Name    string `json:"name"`
		Content string `json:"content"`
	}
	if err := c.ShouldBindJSON(&body); err != nil || strings.TrimSpace(body.Name) == "" || strings.TrimSpace(body.Content) == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "name/content required"})
		return
	}
	owner := userID(c)
	_ = h.db.EnsureKB(c.Request.Context(), owner, rawKB(body.KB), "general")
	// 落库 + 重建索引由入库工作队列统一处理（forceDoc=name 保持笔记身份）。
	job, err := h.enqueueIngest(c.Request.Context(), owner, rawKB(body.KB), scopedKB(c, body.KB), body.Name, "", nil, body.Content)
	if err != nil {
		c.JSON(http.StatusServiceUnavailable, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusAccepted, gin.H{"job_id": job, "name": body.Name})
}

// wikiLinks 从内容抽取所有 [[名称]]（忽略别名）去重，用于维护双链索引。
func wikiLinks(s string) []string {
	seen := map[string]bool{}
	var out []string
	for _, m := range wikiRe.FindAllStringSubmatch(s, -1) {
		n := strings.TrimSpace(m[1])
		if n != "" && !seen[n] {
			seen[n] = true
			out = append(out, n)
		}
	}
	return out
}

var wikiRe = regexp.MustCompile(`\[\[([^\]|]+)(\|[^\]]*)?\]\]`)

// KbVault: GET /api/v1/kb/vault?kb= —— 某知识库的全部原始文档（名+内容），供 Obsidian 式文库浏览。
func (h *Handler) KbVault(c *gin.Context) {
	rows, err := h.db.ListVault(c.Request.Context(), userID(c), rawKB(c.Query("kb")))
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	docs := make([]gin.H, 0, len(rows))
	for _, r := range rows {
		docs = append(docs, gin.H{"id": r.ID, "name": r.Name, "ext": r.Ext, "size": r.Size, "preview": r.Preview})
	}
	c.JSON(http.StatusOK, gin.H{"docs": docs})
}

// KbDoc: GET /api/v1/kb/doc?id= —— 按文件 ID 取单篇全文（按需加载，不在列表里拉全量）。
func (h *Handler) KbDoc(c *gin.Context) {
	d, err := h.db.GetDocByID(c.Request.Context(), userID(c), c.Query("id"))
	if err != nil || d == nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "文档不存在"})
		return
	}
	content := d.Content
	if d.ObjectKey != "" && h.blob.Ready() { // 大文档：从 MinIO 取回正文
		if obj, oerr := h.blob.Get(c.Request.Context(), d.ObjectKey); oerr == nil {
			content = obj
		}
	}
	c.JSON(http.StatusOK, gin.H{"id": d.ID, "name": d.Name, "ext": d.Ext, "content": content, "size": d.Size})
}

// KbDeleteDoc: DELETE /api/v1/kb/doc?id= —— 级联删一份文档：
// 三库(向量/全文/图谱，经 mcp-go kb_delete 按 file_id) + MinIO 原文 + PG 元数据/双链。owner 作用域防越权。
func (h *Handler) KbDeleteDoc(c *gin.Context) {
	owner := userID(c)
	d, err := h.db.GetDocByID(c.Request.Context(), owner, c.Query("id"))
	if err != nil || d == nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "文档不存在"})
		return
	}
	scoped := d.Owner + "/" + d.KB
	// ① 三库按 file_id 级联删（失败不中断，避免半残；记录降级）。
	if res, e := h.bus.CallTool(c.Request.Context(), contract.ToolSubjectGo("kb_delete"),
		&contract.ToolCall{Tool: "kb_delete", Args: map[string]any{"kb": scoped, "file_id": d.ID}}); e != nil || res == nil || !res.OK {
		msg := "kb_delete 无响应"
		if e != nil {
			msg = e.Error()
		} else if res != nil {
			msg = res.Error
		}
		log.Printf("[gateway] 三库删除降级 id=%s: %s", d.ID, msg)
	}
	// ② MinIO 原文。
	if d.ObjectKey != "" && h.blob.Ready() {
		h.blob.Delete(c.Request.Context(), d.ObjectKey)
	}
	// ③ PG 元数据 + 双链。
	if e := h.db.DeleteDocByID(c.Request.Context(), owner, d.KB, d.ID); e != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": e.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"deleted": d.ID, "name": d.Name})
}

// KbLinks: GET /api/v1/kb/links?kb= —— 某库已解析的 [[双链]]（FromID→ToID），供反链/笔记关系图按 ID 渲染。
func (h *Handler) KbLinks(c *gin.Context) {
	rows, err := h.db.ListLinks(c.Request.Context(), userID(c), rawKB(c.Query("kb")))
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	links := make([]gin.H, 0, len(rows))
	for _, l := range rows {
		links = append(links, gin.H{"from": l.FromID, "to": l.ToID})
	}
	c.JSON(http.StatusOK, gin.H{"links": links})
}

// noteName 取文本首个非空行作笔记名（截断 40 字），用于文本入库的文库留存。
func noteName(text string) string {
	for _, line := range strings.Split(text, "\n") {
		line = strings.TrimSpace(line)
		if line != "" {
			r := []rune(line)
			if len(r) > 40 {
				return string(r[:40])
			}
			return line
		}
	}
	return "笔记"
}

// KbIngestFile: POST /api/v1/kb/ingest_file（multipart）—— 文件入库（异步，返回 job_id）。
// 流水线（解析→切块→向量化→写入）的进度经 sundynix.streams.<job_id> 回流，UI 用 SSE 看。
func (h *Handler) KbIngestFile(c *gin.Context) {
	kb := c.PostForm("kb")
	fh, err := c.FormFile("file")
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "file required"})
		return
	}
	f, err := fh.Open()
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	defer f.Close()
	data, err := io.ReadAll(f)
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	_ = h.db.EnsureKB(c.Request.Context(), userID(c), rawKB(kb), "general")
	job, err := h.enqueueIngest(c.Request.Context(), userID(c), rawKB(kb), scopedKB(c, kb), "", fh.Filename, data, "")
	if err != nil {
		c.JSON(http.StatusServiceUnavailable, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusAccepted, gin.H{"job_id": job, "file": fh.Filename})
}

// StartIngestWorkers 启动入库 worker 池：在 JetStream 入库队列上有界并发消费作业（背压、崩溃重投）。
// 返回 drain 供优雅停机时等在途入库跑完（超时未完的由 AckWait 在重启后重投兜底）。
func StartIngestWorkers(ctx context.Context, db *store.Postgres, cache *store.Redis, bus *nats.Bus, blobStore *blob.Store) (func(context.Context), error) {
	h := New(db, cache, bus, blobStore)
	return bus.ConsumeIngestJobs(ctx, h.processIngestJob)
}

// enqueueIngest 把一次入库请求暂存到对象存储(claim-check)，再发布作业到 JetStream 工作队列，
// 立即返回 job_id。暂存让作业消息恒小（不受 max_payload 限制）、崩溃重投只搬引用。
// MinIO 是准生产硬依赖（正文也一律落 MinIO）；未就绪即拒绝入库，不静默退化。
func (h *Handler) enqueueIngest(ctx context.Context, owner, kbName, scoped, forceDoc, filename string, data []byte, rawText string) (string, error) {
	if !h.blob.Ready() {
		return "", errors.New("对象存储未就绪，暂时无法入库")
	}
	job := newJobID()
	stageKey := "ingest-staging/" + job
	isText := filename == ""
	payload := rawText
	if !isText {
		payload = string(data) // 原始文件字节（string 可承载任意字节）
	}
	if err := h.blob.Put(ctx, stageKey, payload); err != nil {
		return "", fmt.Errorf("暂存失败: %w", err)
	}
	jobMsg := &contract.IngestJob{
		JobID: job, Owner: owner, KBName: kbName, Scoped: scoped,
		ForceDoc: forceDoc, Filename: filename, StageKey: stageKey, IsText: isText,
	}
	if err := h.bus.PublishIngestJob(ctx, jobMsg); err != nil {
		h.blob.Delete(ctx, stageKey) // 入队失败 → 清暂存避免泄漏
		return "", fmt.Errorf("入队失败: %w", err)
	}
	return job, nil
}

// processIngestJob 是入库 worker：从暂存取回原料 → 跑入库流水线 → 成功/彻底失败后清暂存。
// 返回非 nil → 队列延迟重投（仅瞬时基建失败且未到最后一次投递时）。幂等：按 doc 先删后写，重跑安全。
func (h *Handler) processIngestJob(ctx context.Context, job *contract.IngestJob, lastAttempt bool) error {
	payload, err := h.blob.Get(ctx, job.StageKey)
	if err != nil {
		// 暂存读不到（已清/对象丢失）——无法恢复，别无限重试，直接终态收尾。
		_ = h.bus.PublishIngest(job.JobID, &contract.IngestEvent{Stage: "失败", Error: "暂存读取失败: " + err.Error()})
		_ = h.bus.CompleteStream(job.JobID)
		return nil
	}
	var data []byte
	var rawText string
	if job.IsText {
		rawText = payload
	} else {
		data = []byte(payload)
	}
	retryable, rerr := h.runIngest(ctx, job.JobID, job.Owner, job.KBName, job.Scoped, job.ForceDoc, job.Filename, data, rawText)
	if rerr != nil && retryable && !lastAttempt {
		return rerr // 瞬时失败 → 保留暂存，延迟重投
	}
	h.blob.Delete(context.Background(), job.StageKey) // 成功或彻底失败 → 清暂存
	return nil
}

// runIngest 跑入库流水线，逐阶段把进度发到 sundynix.streams.<job>，由入库 worker 调用。
// owner+kbName 用于"文库"原文留存；scoped 是 owner/kb 作向量/全文/图谱分区键。
// forceDoc 非空时强制以它为文档名（笔记编辑用，保持笔记身份稳定）。
// filename 非空表示文件入库（先经 mcp-py 解析）；否则用 rawText。
// 返回 (retryable, err)：解析失败=终态(坏输入重试无益)；kb_ingest 基建失败=瞬时可重试。
func (h *Handler) runIngest(ctx context.Context, job, owner, kbName, scoped, forceDoc, filename string, data []byte, rawText string) (retryable bool, err error) {
	emit := func(ev contract.IngestEvent) { _ = h.bus.PublishIngest(job, &ev) }
	time.Sleep(400 * time.Millisecond) // 给 SSE 客户端订阅时间（core NATS 无缓冲）

	text := rawText
	if filename != "" {
		emit(contract.IngestEvent{Stage: "解析", Msg: filename})
		parsed, perr := h.parseFile(ctx, filename, data)
		if perr != nil {
			emit(contract.IngestEvent{Stage: "失败", Error: "解析失败: " + perr.Error()})
			_ = h.bus.CompleteStream(job)
			return false, fmt.Errorf("解析失败: %w", perr) // 终态：坏输入
		}
		emit(contract.IngestEvent{
			Stage:   "解析完成",
			Msg:     fmt.Sprintf("%s · 解析出 %d 字", fileKind(filename), len([]rune(parsed))),
			Preview: head(parsed, 240),
		})
		text = parsed
	}

	// 文库留存原文：编辑指定名 > 文件名 > 文本首行。
	docName := forceDoc
	if docName == "" {
		docName = filename
	}
	if docName == "" {
		docName = noteName(text)
	}
	var fileID string // 文档稳定 ID（雪花），作 Milvus/Bleve/Neo4j 的关联键——重名/重入库不变，删可级联
	if text != "" {
		size := len([]rune(text))
		ext := strings.ToLower(filepath.Ext(filename)) // 笔记/文本入库时 filename 为空 → ext 为空
		inline, objectKey := text, ""
		// 正文一律落对象存储（MinIO），PG 只留元数据+预览+对象键（不分大小，不把正文塞进 PG）。
		// 仅当 MinIO 不可用或写失败时，才回退内联，保证正文不丢。
		if h.blob.Ready() {
			key := owner + "/" + kbName + "/" + docName
			if err := h.blob.Put(ctx, key, text); err == nil {
				inline, objectKey = "", key
			} else {
				log.Printf("[gateway] 正文转 MinIO 失败，回退内联: %v", err)
			}
		} else {
			log.Printf("[gateway] MinIO 未就绪，正文回退内联存 PG（doc=%s）", docName)
		}
		docID, oldKey, err := h.db.SaveDoc(ctx, owner, kbName, docName, ext, inline, objectKey, size, head(text, 500))
		if err != nil {
			log.Printf("[gateway] 文件入库失败: %v", err)
		} else if docID != "" {
			fileID = docID // 下游三库用它作关联键
			// 孤儿 GC：重名覆盖后旧对象键若已不用（转内联或换键），从 MinIO 删除，避免泄漏。
			if oldKey != "" && oldKey != objectKey && h.blob.Ready() {
				h.blob.Delete(ctx, oldKey)
				log.Printf("[gateway] 清理被覆盖的 MinIO 孤儿对象: %s", oldKey)
			}
			_ = h.db.ReplaceDocLinks(ctx, owner, kbName, docID, wikiLinks(text)) // 以本文件 ID 维护出链
			_ = h.db.ResolveInboundLinks(ctx, owner, kbName, docName, docID)      // 回填指向本文件的悬空链接
		}
	}

	// 调 mcp-go kb_ingest：doc 传**稳定 file_id**（非展示名），作下游三库关联键，使重名/重入库幂等、删可级联。
	// SaveDoc 失败兜底退回用名字做键（仍能入库，只是失去 file_id 稳定性）。
	docKey := fileID
	if docKey == "" {
		docKey = docName
	}
	res, cerr := h.bus.CallTool(ctx, contract.ToolSubjectGo("kb_ingest"),
		&contract.ToolCall{Tool: "kb_ingest", Args: map[string]any{"kb": scoped, "doc": docKey, "text": text, "job_id": job}})
	if cerr != nil || res == nil || !res.OK {
		msg := "kb_ingest 失败"
		if cerr != nil {
			msg = cerr.Error()
		} else if res != nil {
			msg = res.Error
		}
		emit(contract.IngestEvent{Stage: "失败", Error: msg})
		_ = h.bus.CompleteStream(job)
		return true, errors.New(msg) // 瞬时：mcp-go/embedding/Milvus 抖动，可重试
	}
	return false, nil
}

// KbIngestStream: GET /api/v1/kb/ingest/:id/stream —— SSE 实时推送入库进度事件。
func (h *Handler) KbIngestStream(c *gin.Context) {
	job := c.Param("id")
	c.Writer.Header().Set("Content-Type", "text/event-stream")
	c.Writer.Header().Set("Cache-Control", "no-cache")
	c.Writer.Header().Set("Connection", "keep-alive")

	events := make(chan []byte, 64)
	done := make(chan struct{})
	unsub, err := h.bus.SubscribeTokens(job,
		func(ev []byte) {
			select {
			case events <- ev:
			default:
			}
		},
		func() { close(done) },
	)
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	defer func() { _ = unsub() }()

	c.Stream(func(w io.Writer) bool {
		select {
		case ev := <-events:
			c.SSEvent("progress", string(ev))
			return true
		case <-done:
			c.SSEvent("done", job)
			return false
		case <-c.Request.Context().Done():
			return false
		}
	})
}

func newJobID() string {
	var b [8]byte
	_, _ = rand.Read(b[:])
	return "ingest_" + hex.EncodeToString(b[:])
}

// itoa 简易整数转字符串（避免引入 strconv）。
func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	neg := n < 0
	if neg {
		n = -n
	}
	var b []byte
	for n > 0 {
		b = append([]byte{byte('0' + n%10)}, b...)
		n /= 10
	}
	if neg {
		b = append([]byte{'-'}, b...)
	}
	return string(b)
}

// head 取文本前 n 个字符（按 rune），用于解析预览。
func head(s string, n int) string {
	s = strings.TrimSpace(s)
	r := []rune(s)
	if len(r) <= n {
		return s
	}
	return string(r[:n]) + "…"
}

// fileKind 由扩展名给出可读的文件类型标签。
func fileKind(filename string) string {
	switch strings.ToLower(filepath.Ext(filename)) {
	case ".docx":
		return "Word 文档"
	case ".xlsx", ".xls":
		return "Excel 表格"
	case ".pdf":
		return "PDF"
	case ".csv":
		return "CSV"
	case ".md", ".markdown":
		return "Markdown"
	default:
		return "文本"
	}
}

// parseFile 把文件字节转为纯文本：文本类直读，其余经 mcp-py parse_document（算法层）。
func (h *Handler) parseFile(ctx context.Context, filename string, data []byte) (string, error) {
	switch strings.ToLower(filepath.Ext(filename)) {
	case ".txt", ".md", ".markdown", ".text":
		return string(data), nil
	}
	res, err := h.bus.CallTool(ctx, contract.ToolSubjectPy("parse_document"),
		&contract.ToolCall{Tool: "parse_document", Args: map[string]any{
			"filename": filename, "content_b64": base64.StdEncoding.EncodeToString(data),
		}})
	if err != nil {
		return "", err
	}
	if res == nil || !res.OK {
		if res != nil && res.Error != "" {
			return "", errors.New(res.Error)
		}
		return "", errors.New("parse_document 无响应（mcp-py 未运行？）")
	}
	return res.Content, nil
}

// KbSearch: POST /api/v1/kb/search —— 检索台：查某知识库，返回带分数的命中（→ mcp-go kb_search）。
func (h *Handler) KbSearch(c *gin.Context) {
	var body struct {
		KB   string `json:"kb"`
		Q    string `json:"q"`
		TopK int    `json:"topK"`
		Mode string `json:"mode"` // 空=生产混合(含rerank)；vector/fulltext/graph/hybrid=评测单路/纯融合
	}
	if err := c.ShouldBindJSON(&body); err != nil || body.Q == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "q required"})
		return
	}
	args := map[string]any{"kb": scopedKB(c, body.KB), "q": body.Q}
	if body.TopK > 0 {
		args["topK"] = body.TopK
	}
	if body.Mode != "" {
		args["mode"] = body.Mode
	}
	res, err := h.bus.CallTool(c.Request.Context(), contract.ToolSubjectGo("kb_search"),
		&contract.ToolCall{Tool: "kb_search", Args: args})
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	if !res.OK {
		c.JSON(http.StatusUnprocessableEntity, gin.H{"error": res.Error})
		return
	}
	var hits []map[string]any
	_ = json.Unmarshal([]byte(res.Content), &hits)
	c.JSON(http.StatusOK, gin.H{"hits": hits})
}

// KbGraph: GET /api/v1/kb/graph?kb= —— 某知识库的图谱三元组（→ mcp-go kb_graph，Neo4j）。
func (h *Handler) KbGraph(c *gin.Context) {
	res, err := h.bus.CallTool(c.Request.Context(), contract.ToolSubjectGo("kb_graph"),
		&contract.ToolCall{Tool: "kb_graph", Args: map[string]any{"kb": scopedKB(c, c.Query("kb")), "limit": 100}})
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	var triples []map[string]any
	_ = json.Unmarshal([]byte(res.Content), &triples)
	c.JSON(http.StatusOK, gin.H{"triples": triples})
}