FranklinWH-Toolbox/service/cmd/mcp/tools/data/mongo.go

package data

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo/options"

	"fwh-toolbox-server/cmd/mcp/tools/helper"
	"fwh-toolbox-server/database"
)

// mongoAllowedCollections 允许查询的 collection 白名单
var mongoAllowedCollections = []string{
	// 站点信息管理
	"SiteInfoManageData",
	"SiteInfoManageRecord",
	// SIM卡管理
	"SimCardManage",
	"SimCardMonthlyData",
	// 批量下载
	"BatchDownloadConfig",
	// 站点分析
	"SiteAnalysisConfig",
	// 统计
	"StatisticsShutdown",
	"StatisticsAlarm",
	// 站点监控
	"SiteMonitorConfig",
	// 自动升级
	"AutoUpdateConfigOption",
	"AutoUpdateConfig",
	// 注意：UserInfo（内部用户身份数据）与 AiChatSessions（历史对话自由文本，
	// 字段级脱敏无效）已从白名单移除，禁止通过 data_mongo_* 直查以防隐私泄露。
}

func registerMongoTools(s *server.MCPServer) {
	s.AddTool(
		mcp.NewTool("data_mongo_find",
			mcp.WithDescription("直接查询 MongoDB 集合数据（只读）。用于 service 层无法满足的灵活查询场景。"),
			mcp.WithString("collection", mcp.Required(), mcp.Description(
				fmt.Sprintf("集合名，可选值: %v", mongoAllowedCollections),
			)),
			mcp.WithString("filter", mcp.Description("查询条件，JSON 格式的 MongoDB filter（如 {\"sn\": \"xxx\"}），不传则查全部")),
			mcp.WithArray("fields", mcp.WithStringItems(),
				mcp.Description("字段投影：要返回的字段名数组（如 [\"sn\",\"status\"]），留空返回全部字段。用于大数据量场景下减少返回体积、节省 token。")),
			mcp.WithNumber("limit", mcp.Description("返回数量限制，默认10，最大100")),
			mcp.WithNumber("skip", mcp.Description("跳过数量（分页偏移），默认0")),
			mcp.WithString("sort", mcp.Description("排序条件，JSON 格式（如 {\"createdAt\": -1}）")),
		),
		handleMongoFind,
	)

	s.AddTool(
		mcp.NewTool("data_mongo_count",
			mcp.WithDescription("统计 MongoDB 集合中符合条件的文档数量"),
			mcp.WithString("collection", mcp.Required(), mcp.Description(
				fmt.Sprintf("集合名，可选值: %v", mongoAllowedCollections),
			)),
			mcp.WithString("filter", mcp.Description("查询条件，JSON 格式的 MongoDB filter，不传则统计全部")),
		),
		handleMongoCount,
	)

	s.AddTool(
		mcp.NewTool("data_mongo_export",
			mcp.WithDescription("全量导出 MongoDB 集合数据到本地文件（只读）。用于大数据量导出场景：不受 find 的 100 条上限限制，游标流式拉取符合条件的全部文档，边脱敏边写入 NDJSON 文件（每行一个 JSON 对象），内存占用为常数级。仅返回 localPath、count 与结构化摘要（不返回数据正文，避免占用上下文）。为防止误导出全表，无 filter 时若总数超过上限会被拒绝。"),
			mcp.WithString("collection", mcp.Required(), mcp.Description(
				fmt.Sprintf("集合名，可选值: %v", mongoAllowedCollections),
			)),
			mcp.WithString("filter", mcp.Description("查询条件，JSON 格式的 MongoDB filter，不传则导出全部（受上限保护）")),
			mcp.WithArray("fields", mcp.WithStringItems(),
				mcp.Description("字段投影：只导出这些字段（如 [\"sn\",\"status\"]），留空导出全部字段。")),
			mcp.WithString("sort", mcp.Description("排序条件，JSON 格式（如 {\"createdAt\": -1}）")),
			mcp.WithNumber("max_docs", mcp.Description(fmt.Sprintf("最大导出条数上限，默认 %d，硬上限 %d。超过则拒绝，请缩小 filter。", mongoExportDefaultMax, mongoExportHardMax))),
		),
		handleMongoExport,
	)

	s.AddTool(
		mcp.NewTool("data_mongo_aggregate",
			mcp.WithDescription("对 MongoDB 集合执行聚合管道（aggregation pipeline，只读）并将结果流式导出到本地文件。"+
				"支持 $match/$project/$group/$addFields/$lookup（跨集合关联/JOIN）/$sort/$limit/$unwind/$facet 等只读阶段，"+
				"可在数据库侧一次性完成过滤+关联+聚合，避免多次取数与 Agent 侧拼接。"+
				"结果边脱敏边写入 NDJSON 文件（每行一个 JSON 对象），仅返回 localPath+count+summary（不返回数据正文）。"+
				"安全约束：collection 与 $lookup/$unionWith/$graphLookup 引用的集合都必须在白名单内；"+
				"禁止 $out/$merge/$function/$accumulator/$where 等写入或执行代码的阶段。"),
			mcp.WithString("collection", mcp.Required(), mcp.Description(
				fmt.Sprintf("聚合起点集合名，可选值: %v", mongoAllowedCollections),
			)),
			mcp.WithString("pipeline", mcp.Required(), mcp.Description(
				"聚合管道，JSON 数组字符串，每个元素是一个聚合阶段。"+
					"例：[{\"$match\":{\"province\":\"Puerto Rico\"}},{\"$group\":{\"_id\":\"$province\",\"n\":{\"$sum\":1}}}]")),
			mcp.WithNumber("max_docs", mcp.Description(fmt.Sprintf(
				"最大输出条数上限，默认 %d，硬上限 %d。聚合结果超过上限会被截断并在返回中标注。", mongoAggregateDefaultMax, mongoExportHardMax))),
		),
		handleMongoAggregate,
	)
}

// mongoAggregateDefaultMax 聚合输出默认条数上限（软保护，可由 max_docs 覆盖）。
const mongoAggregateDefaultMax = 50000

// mongoAllowedStages 聚合管道允许的只读阶段白名单。
// 明确排除写入/执行代码的阶段：$out、$merge（写集合）、$function、$accumulator、$where（执行 JS）。
var mongoAllowedStages = map[string]bool{
	"$match":         true,
	"$project":       true,
	"$addFields":     true,
	"$set":           true,
	"$unset":         true,
	"$group":         true,
	"$sort":          true,
	"$limit":         true,
	"$skip":          true,
	"$count":         true,
	"$unwind":        true,
	"$lookup":        true,
	"$graphLookup":   true,
	"$unionWith":     true,
	"$facet":         true,
	"$bucket":        true,
	"$bucketAuto":    true,
	"$sortByCount":   true,
	"$replaceRoot":   true,
	"$replaceWith":   true,
	"$sample":        true,
	"$redact":        true,
	"$setWindowFields": true,
}

// mongoExportDefaultMax 全量导出默认上限（防误导出全表的软保护，可由 max_docs 覆盖）
const mongoExportDefaultMax = 50000

// mongoExportHardMax 全量导出硬上限。流式导出内存为常数级，此上限仅用于防止
// 误导出超大规模数据，而非内存保护，故设置较高。
const mongoExportHardMax = 5000000

func handleMongoFind(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	collection, err := req.RequireString("collection")
	if err != nil {
		return helper.ParamError("collection"), nil
	}

	if !isAllowedCollection(collection) {
		return helper.ParamInvalidError("collection", fmt.Sprintf("不允许查询此集合，可选: %v", mongoAllowedCollections)), nil
	}

	// 解析 filter
	filterStr := req.GetString("filter", "{}")
	var filter bson.M
	if err := json.Unmarshal([]byte(filterStr), &filter); err != nil {
		return helper.ParamInvalidError("filter", fmt.Sprintf("JSON 解析失败: %v", err)), nil
	}

	// limit
	limit := int64(req.GetFloat("limit", 10))
	if limit > 100 {
		limit = 100
	}
	if limit < 1 {
		limit = 10
	}

	// skip
	skip := int64(req.GetFloat("skip", 0))
	if skip < 0 {
		skip = 0
	}

	// sort
	var sort bson.M
	sortStr := req.GetString("sort", "")
	if sortStr != "" {
		if err := json.Unmarshal([]byte(sortStr), &sort); err != nil {
			return helper.ParamInvalidError("sort", fmt.Sprintf("JSON 解析失败: %v", err)), nil
		}
	}

	// 查询
	coll := database.GetCollection(collection)
	findOpts := options.Find().SetLimit(limit).SetSkip(skip)
	if sort != nil {
		findOpts.SetSort(sort)
	}

	// fields 字段投影：留空返回全部字段；指定则只返回这些字段（源头减量，节省 token）
	if fields, ok := helper.GetStringSlice(req, "fields"); ok && len(fields) > 0 {
		projection := bson.M{}
		for _, f := range fields {
			projection[f] = 1
		}
		findOpts.SetProjection(projection)
	}

	cursor, err := coll.Find(ctx, filter, findOpts)
	if err != nil {
		return helper.ServiceError("MongoDB查询", err), nil
	}
	defer cursor.Close(ctx)

	var results []bson.M
	if err := cursor.All(ctx, &results); err != nil {
		return helper.ServiceError("读取MongoDB结果", err), nil
	}

	// total：符合 filter 的文档总数（不受 limit/skip 影响），便于大数据量遍历时判断进度
	total, err := coll.CountDocuments(ctx, filter)
	if err != nil {
		return helper.ServiceError("统计MongoDB总数", err), nil
	}

	returned := int64(len(results))
	hasMore := skip+returned < total

	// 结构化返回：附带 total/returned/skip/limit/hasMore，便于 AI 判断是否翻完、如何翻页
	return helper.JSONResultOrFileAuto(map[string]any{
		"total":    total,
		"returned": returned,
		"skip":     skip,
		"limit":    limit,
		"hasMore":  hasMore,
		"items":    results,
	}), nil
}

func handleMongoCount(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	collection, err := req.RequireString("collection")
	if err != nil {
		return helper.ParamError("collection"), nil
	}

	if !isAllowedCollection(collection) {
		return helper.ParamInvalidError("collection", fmt.Sprintf("不允许查询此集合，可选: %v", mongoAllowedCollections)), nil
	}

	filterStr := req.GetString("filter", "{}")
	var filter bson.M
	if err := json.Unmarshal([]byte(filterStr), &filter); err != nil {
		return helper.ParamInvalidError("filter", fmt.Sprintf("JSON 解析失败: %v", err)), nil
	}

	coll := database.GetCollection(collection)
	count, err := coll.CountDocuments(ctx, filter)
	if err != nil {
		return helper.ServiceError("MongoDB统计", err), nil
	}

	return helper.JSONResultOrFileAuto(map[string]any{
		"collection": collection,
		"count":      count,
	}), nil
}

func handleMongoExport(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	collection, err := req.RequireString("collection")
	if err != nil {
		return helper.ParamError("collection"), nil
	}
	if !isAllowedCollection(collection) {
		return helper.ParamInvalidError("collection", fmt.Sprintf("不允许查询此集合，可选: %v", mongoAllowedCollections)), nil
	}

	// 解析 filter
	filterStr := req.GetString("filter", "{}")
	var filter bson.M
	if err := json.Unmarshal([]byte(filterStr), &filter); err != nil {
		return helper.ParamInvalidError("filter", fmt.Sprintf("JSON 解析失败: %v", err)), nil
	}
	hasFilter := len(filter) > 0

	// max_docs 上限
	maxDocs := int64(req.GetFloat("max_docs", float64(mongoExportDefaultMax)))
	if maxDocs < 1 {
		maxDocs = mongoExportDefaultMax
	}
	if maxDocs > mongoExportHardMax {
		maxDocs = mongoExportHardMax
	}

	coll := database.GetCollection(collection)

	// 先统计总量，做防误导出保护
	total, err := coll.CountDocuments(ctx, filter)
	if err != nil {
		return helper.ServiceError("统计导出总数", err), nil
	}
	if total == 0 {
		return helper.EmptyResult("符合条件的文档为 0，无需导出"), nil
	}
	if total > maxDocs {
		hint := "请缩小 filter 或调高 max_docs"
		if !hasFilter {
			hint = "未提供 filter 且数据量过大，请提供 filter 缩小范围，或显式调高 max_docs"
		}
		return helper.ParamInvalidError("max_docs", fmt.Sprintf(
			"符合条件的文档数 %d 超过上限 %d，%s（硬上限 %d）", total, maxDocs, hint, mongoExportHardMax)), nil
	}

	// 组装查询选项
	findOpts := options.Find()
	if sortStr := req.GetString("sort", ""); sortStr != "" {
		var sort bson.M
		if err := json.Unmarshal([]byte(sortStr), &sort); err != nil {
			return helper.ParamInvalidError("sort", fmt.Sprintf("JSON 解析失败: %v", err)), nil
		}
		findOpts.SetSort(sort)
	}
	if fields, ok := helper.GetStringSlice(req, "fields"); ok && len(fields) > 0 {
		projection := bson.M{}
		for _, f := range fields {
			projection[f] = 1
		}
		findOpts.SetProjection(projection)
	}

	// 游标流式导出：逐条 decode → 脱敏 → 写 NDJSON，内存常数级
	cursor, err := coll.Find(ctx, filter, findOpts)
	if err != nil {
		return helper.ServiceError("MongoDB导出查询", err), nil
	}
	defer cursor.Close(ctx)

	res := helper.ExportStreamNDJSON(func() (any, bool, error) {
		if !cursor.Next(ctx) {
			// Next 返回 false：可能是正常结束，也可能是出错
			return nil, false, cursor.Err()
		}
		var doc bson.M
		if err := cursor.Decode(&doc); err != nil {
			return nil, false, err
		}
		return doc, true, nil
	}, "mongo_export_"+collection)
	return res, nil
}

func handleMongoAggregate(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	collection, err := req.RequireString("collection")
	if err != nil {
		return helper.ParamError("collection"), nil
	}
	if !isAllowedCollection(collection) {
		return helper.ParamInvalidError("collection", fmt.Sprintf("不允许查询此集合，可选: %v", mongoAllowedCollections)), nil
	}

	// 解析 pipeline（JSON 数组）
	pipelineStr, err := req.RequireString("pipeline")
	if err != nil {
		return helper.ParamError("pipeline"), nil
	}
	var pipeline []bson.M
	if err := json.Unmarshal([]byte(pipelineStr), &pipeline); err != nil {
		return helper.ParamInvalidError("pipeline", fmt.Sprintf("需要 JSON 数组，每个元素是一个聚合阶段: %v", err)), nil
	}
	if len(pipeline) == 0 {
		return helper.ParamInvalidError("pipeline", "pipeline 不能为空"), nil
	}

	// 安全校验：阶段白名单 + 引用集合白名单（递归）
	if verr := validatePipeline(pipeline); verr != "" {
		return helper.ParamInvalidError("pipeline", verr), nil
	}

	// max_docs 输出上限
	maxDocs := int64(req.GetFloat("max_docs", float64(mongoAggregateDefaultMax)))
	if maxDocs < 1 {
		maxDocs = mongoAggregateDefaultMax
	}
	if maxDocs > mongoExportHardMax {
		maxDocs = mongoExportHardMax
	}

	// 聚合可能较重：设置超时 + allowDiskUse 防止内存溢出。
	aggCtx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	coll := database.GetCollection(collection)
	allowDisk := true
	cursor, err := coll.Aggregate(aggCtx, pipeline, &options.AggregateOptions{AllowDiskUse: &allowDisk})
	if err != nil {
		return helper.ServiceError("MongoDB聚合", err), nil
	}
	defer cursor.Close(aggCtx)

	// 流式：逐条 decode → 脱敏 → 写 NDJSON，内存常数级；达到 maxDocs 即停止。
	var emitted int64
	res := helper.ExportStreamNDJSON(func() (any, bool, error) {
		if emitted >= maxDocs {
			return nil, false, nil // 达到上限，正常结束（截断）
		}
		if !cursor.Next(aggCtx) {
			return nil, false, cursor.Err()
		}
		var doc bson.M
		if err := cursor.Decode(&doc); err != nil {
			return nil, false, err
		}
		emitted++
		return doc, true, nil
	}, "mongo_aggregate_"+collection)
	return res, nil
}

// validatePipeline 对聚合管道做安全校验，返回空字符串表示通过，否则返回错误说明。
//   - 每个阶段的顶层操作符必须在 mongoAllowedStages 白名单内（排除 $out/$merge/$function 等）。
//   - 递归收集 $lookup/$graphLookup/$unionWith 引用的集合，全部必须在集合白名单内，
//     防止借跨集合阶段绕过白名单读取被禁集合。
func validatePipeline(pipeline []bson.M) string {
	for i, stage := range pipeline {
		if len(stage) != 1 {
			return fmt.Sprintf("第 %d 个阶段必须恰好包含一个聚合操作符", i+1)
		}
		for op, spec := range stage {
			if !mongoAllowedStages[op] {
				return fmt.Sprintf("第 %d 个阶段使用了不允许的操作符 %q（仅允许只读阶段，禁止 $out/$merge/$function 等）", i+1, op)
			}
			if refs := collectStageCollectionRefs(op, spec); len(refs) > 0 {
				for _, c := range refs {
					if !isAllowedCollection(c) {
						return fmt.Sprintf("第 %d 个阶段 %s 引用了不在白名单内的集合 %q，可选: %v", i+1, op, c, mongoAllowedCollections)
					}
				}
			}
			// $unionWith 可能携带嵌套 pipeline，需递归校验
			if op == "$unionWith" {
				if sub := extractUnionWithPipeline(spec); sub != nil {
					if verr := validatePipeline(sub); verr != "" {
						return fmt.Sprintf("第 %d 个阶段 $unionWith 的嵌套 pipeline 校验失败: %s", i+1, verr)
					}
				}
			}
			// $lookup 的子 pipeline 形式（带 pipeline 字段）也需递归校验
			if op == "$lookup" {
				if sub := extractLookupPipeline(spec); sub != nil {
					if verr := validatePipeline(sub); verr != "" {
						return fmt.Sprintf("第 %d 个阶段 $lookup 的嵌套 pipeline 校验失败: %s", i+1, verr)
					}
				}
			}
		}
	}
	return ""
}

// collectStageCollectionRefs 提取某个阶段引用的外部集合名（用于白名单校验）。
//   - $lookup:      spec.from（字符串）
//   - $graphLookup: spec.from（字符串）
//   - $unionWith:   字符串本身，或 spec.coll（字符串）
func collectStageCollectionRefs(op string, spec any) []string {
	var refs []string
	switch op {
	case "$lookup", "$graphLookup":
		if m, ok := toMap(spec); ok {
			if from, ok := m["from"].(string); ok && from != "" {
				refs = append(refs, from)
			}
		}
	case "$unionWith":
		if s, ok := spec.(string); ok && s != "" {
			refs = append(refs, s)
		} else if m, ok := toMap(spec); ok {
			if coll, ok := m["coll"].(string); ok && coll != "" {
				refs = append(refs, coll)
			}
		}
	}
	return refs
}

// extractUnionWithPipeline 提取 $unionWith 的嵌套 pipeline（若有）。
func extractUnionWithPipeline(spec any) []bson.M {
	m, ok := toMap(spec)
	if !ok {
		return nil
	}
	return toPipeline(m["pipeline"])
}

// extractLookupPipeline 提取 $lookup 的子 pipeline（若有）。
func extractLookupPipeline(spec any) []bson.M {
	m, ok := toMap(spec)
	if !ok {
		return nil
	}
	return toPipeline(m["pipeline"])
}

// toMap 将 bson.M / map[string]any 统一转为 map[string]any。
func toMap(v any) (map[string]any, bool) {
	switch t := v.(type) {
	case bson.M:
		return map[string]any(t), true
	case map[string]any:
		return t, true
	default:
		return nil, false
	}
}

// toPipeline 将任意值尽力转换为 []bson.M（聚合子管道）。
func toPipeline(v any) []bson.M {
	arr, ok := v.([]any)
	if !ok {
		return nil
	}
	out := make([]bson.M, 0, len(arr))
	for _, e := range arr {
		if m, ok := toMap(e); ok {
			out = append(out, bson.M(m))
		}
	}
	return out
}

func isAllowedCollection(name string) bool {
	for _, c := range mongoAllowedCollections {
		if c == name {
			return true
		}
	}
	return false
}