"""MongoDB 审核规则引擎 — 42+ 条 MongoDB 特有审核规则。

支持:
- 集合管理规则 (Collection Management)
- 查询安全规则 (Query Safety)
- 聚合管道规则 (Aggregation Pipeline)
- 写入安全规则 (Write Safety)
- 安全规则 (Security)
- 性能规则 (Performance)

使用方式:
    engine = SQLAuditEngineV2(db_session)
    report = await engine.audit(
        '{"find": "users", "filter": {"age": {"$gt": 18}}}',
        db_type="mongodb"
    )
"""

from __future__ import annotations

import json
import re
from typing import Any, Dict, List, Optional, Tuple

from app.services.audit_engine_v2 import (
    register_rule,
    AuditRuleV2,
    AuditFindingV2,
    AuditReportV2,
    Severity,
    RuleCategory,
)


# ──────────────────────────────────────────────
# JSON / BSON 辅助解析函数
# ──────────────────────────────────────────────


def _try_parse_json(sql_clean: str) -> Tuple[bool, Dict[str, Any]]:
    """尝试解析 MongoDB 查询为 JSON 对象。

    Returns:
        (success, parsed_dict) — 解析成功返回 (True, dict)，失败返回 (False, {})
    """
    # 去除首尾空白
    text = sql_clean.strip()
    if not text:
        return False, {}
    # 如果文本被单引号包裹，去除外层引号
    if text.startswith("'") and text.endswith("'"):
        text = text[1:-1]
    if text.startswith('"') and text.endswith('"') and len(text) > 1:
        inner = text[1:-1]
        # 检查内部是否是有效的 JSON（避免普通字符串被误剥）
        if inner.strip().startswith("{"):
            text = inner
    # 去除末尾可能存在的分号
    text = text.rstrip(";").strip()
    if not text:
        return False, {}
    try:
        parsed = json.loads(text)
        if isinstance(parsed, dict):
            return True, parsed
        return False, {}
    except (json.JSONDecodeError, ValueError):
        return False, {}


def _is_json_mode(sql_clean: str) -> bool:
    """检查是否为 JSON 格式的 MongoDB 查询。

    MongoDB 查询通常是 JSON 格式（以 { 开头），
    而 shell 命令如 db.collection.find() 以 db. 或 javascript 开头。
    """
    text = sql_clean.strip()
    if not text:
        return False
    # 标准的 JSON 查询格式 {"find": "collection", ...}
    if text.startswith("{"):
        return True
    # 或者 shell 命令中带有 { } 参数
    if "{" in text and "}" in text:
        return True
    return False


def _is_shell_command(sql_clean: str) -> bool:
    """检查是否为 MongoDB shell 命令（如 db.collection.find()）。"""
    text = sql_clean.strip()
    return text.startswith("db.") or text.startswith("sh.")


def _extract_collection_name(sql_clean: str) -> Optional[str]:
    """从 MongoDB 查询中提取集合名称。

    支持格式:
    - JSON: {"find": "collection_name", ...}
    - Shell: db.collection_name.find(...)
    - Shell: db.getCollection("name").find(...)
    """
    # 尝试 JSON 格式
    success, parsed = _try_parse_json(sql_clean)
    if success:
        # {"find": "users", ...}
        if "find" in parsed and isinstance(parsed["find"], str):
            return parsed["find"]
        # {"insert": "users", ...}
        for op in ("insert", "update", "delete", "aggregate", "count", "distinct",
                    "findOne", "insertOne", "updateOne", "deleteOne",
                    "insertMany", "updateMany", "deleteMany"):
            if op in parsed and isinstance(parsed[op], str):
                return parsed[op]
        return None

    # Shell 格式: db.collection_name.operation()
    text = sql_clean.strip()
    # db.getCollection("name").method(...)
    gc_match = re.search(r'\bgetCollection\s*\(\s*["\'](\w+)["\']\s*\)', text)
    if gc_match:
        return gc_match.group(1)
    # db.name.method(...)
    shell_match = re.search(r'\bdb\.(\w+)\.', text)
    if shell_match:
        return shell_match.group(1)
    return None


def _has_where_operator(sql_clean: str) -> bool:
    """检查是否使用了 $where 运算符（注入风险）。"""
    return "\"$where\"" in sql_clean or "'$where'" in sql_clean


def _has_function_operator(sql_clean: str) -> bool:
    """检查是否使用了 $function 运算符。"""
    return "\"$function\"" in sql_clean or "'$function'" in sql_clean


def _count_json_keys(obj: Any) -> int:
    """递归计算 JSON 对象中的键总数（估算文档大小）。"""
    if isinstance(obj, dict):
        return sum(1 + _count_json_keys(v) for k, v in obj.items())
    elif isinstance(obj, list):
        return sum(_count_json_keys(item) for item in obj)
    return 1


def _get_nesting_depth(obj: Any, current_depth: int = 0) -> int:
    """计算 JSON 对象的最大嵌套深度。"""
    if isinstance(obj, dict):
        if not obj:
            return current_depth
        return max(_get_nesting_depth(v, current_depth + 1) for v in obj.values())
    elif isinstance(obj, list):
        if not obj:
            return current_depth
        return max(_get_nesting_depth(item, current_depth + 1) for item in obj)
    return current_depth


def _find_all_values(obj: Any, key: str) -> List[Any]:
    """递归查找 JSON 对象中所有指定 key 的值。"""
    results: List[Any] = []
    if isinstance(obj, dict):
        for k, v in obj.items():
            if k == key:
                results.append(v)
            results.extend(_find_all_values(v, key))
    elif isinstance(obj, list):
        for item in obj:
            results.extend(_find_all_values(item, key))
    return results


def _has_key_recursive(obj: Any, key: str) -> bool:
    """递归检查 JSON 对象是否包含指定 key。"""
    if isinstance(obj, dict):
        if key in obj:
            return True
        return any(_has_key_recursive(v, key) for v in obj.values())
    elif isinstance(obj, list):
        return any(_has_key_recursive(item, key) for item in obj)
    return False


def _get_json_string_length(sql_clean: str) -> int:
    """估算 JSON 字符串的字节长度（用于文档大小估算）。"""
    return len(sql_clean.encode("utf-8"))


def _count_pipeline_stages(sql_clean: str) -> int:
    """计算聚合管道中的阶段数量。"""
    success, parsed = _try_parse_json(sql_clean)
    if success and "pipeline" in parsed and isinstance(parsed["pipeline"], list):
        return len(parsed["pipeline"])
    # shell 格式中粗略统计 $ 开头的 stage
    stages = re.findall(r'\{\\s*"\\$\w+"', sql_clean)
    return len(stages)


def _is_snake_case(name: str) -> bool:
    """检查是否为 snake_case 命名。"""
    if not name:
        return False
    return bool(re.match(r"^[a-z][a-z0-9_]*$", name))


# ═══════════════════════════════════════════════
# 集合管理规则 (8条) — MONGO-COL-001 ~ 008
# ═══════════════════════════════════════════════

@register_rule(AuditRuleV2(
    code="MONGO-COL-001",
    name="集合命名规范(snake_case)",
    category=RuleCategory.DDL,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="MongoDB 集合名应使用小写下划线命名(snake_case)",
    config={"pattern": "^[a-z][a-z0-9_]*$"},
    message_template="集合名 '{name}' 不符合 snake_case 命名规范，应使用小写下划线格式",
))
def check_mongo_col_001(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """集合命名规范: snake_case。"""
    if not _is_json_mode(sql_clean):
        return
    collection_name = _extract_collection_name(sql_clean)
    if not collection_name:
        return
    if not _is_snake_case(collection_name):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(name=collection_name),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-COL-002",
    name="集合名最大长度限制",
    category=RuleCategory.DDL,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="MongoDB 集合名最大长度为 64 字节(不含前缀)",
    config={"max_name_length": 64},
    message_template="集合名 '{name}' 长度为 {length} 字节，超过最大限制 {max} 字节",
))
def check_mongo_col_002(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """集合名最大长度检查: 64 字节。"""
    if not _is_json_mode(sql_clean):
        return
    collection_name = _extract_collection_name(sql_clean)
    if not collection_name:
        return
    max_len = rule.config.get("max_name_length", 64)
    name_bytes = len(collection_name.encode("utf-8"))
    if name_bytes > max_len:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(name=collection_name, length=name_bytes, max=max_len),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-COL-003",
    name="禁止删除集合(DROP)",
    category=RuleCategory.DDL,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="生产环境禁止执行删除集合操作，数据将永久丢失",
    config={"dangerous_ops": ["drop"]},
    message_template="禁止执行删除集合操作: {op}，数据将永久丢失",
))
def check_mongo_col_003(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止删除集合。"""
    if not _is_json_mode(sql_clean):
        return
    success, parsed = _try_parse_json(sql_clean)
    if not success:
        # 检查 shell 命令
        if re.search(r'\bdb\.\w+\.drop\s*\(', sql_clean):
            report.findings.append(AuditFindingV2(
                rule_code=rule.code,
                rule_name=rule.name,
                level=rule.severity,
                message=rule.message_template.format(op="db.collection.drop()"),
            ))
        return
    dangerous = rule.config.get("dangerous_ops", ["drop"])
    for op in dangerous:
        if op in parsed:
            report.findings.append(AuditFindingV2(
                rule_code=rule.code,
                rule_name=rule.name,
                level=rule.severity,
                message=rule.message_template.format(op=op),
            ))
            return


@register_rule(AuditRuleV2(
    code="MONGO-COL-004",
    name="禁止删除数据库(DROP DATABASE)",
    category=RuleCategory.DDL,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="禁止执行 dropDatabase 操作，将删除整个数据库",
    message_template="禁止执行 dropDatabase 操作，将删除整个数据库及其所有集合",
))
def check_mongo_col_004(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止删除数据库。"""
    text = sql_clean.strip()
    # 检查 JSON 格式 {"dropDatabase": 1}
    success, parsed = _try_parse_json(text)
    if success and "dropDatabase" in parsed:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))
        return
    # 检查 shell 格式
    if re.search(r'\bdb\.dropDatabase\s*\(', text):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


@register_rule(AuditRuleV2(
    code="MONGO-COL-005",
    name="禁止删除索引(DROP INDEX)",
    category=RuleCategory.DDL,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="禁止执行删除索引操作，删除后查询性能可能急剧下降",
    message_template="禁止执行删除索引操作: {op}，删除索引可能导致查询性能急剧下降",
))
def check_mongo_col_005(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止删除索引。"""
    text = sql_clean.strip()
    # JSON 格式 {"dropIndexes": "collection", "index": "name"}
    success, parsed = _try_parse_json(text)
    if success and "dropIndexes" in parsed:
        idx_name = parsed.get("index", "*")
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(op=f'dropIndexes: {idx_name}'),
        ))
        return
    # shell 格式: db.collection.dropIndex() / dropIndexes()
    if re.search(r'\b(dropIndex|dropIndexes)\s*\(', text):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(op="dropIndex(s)"),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-COL-006",
    name="集合文档数上限建议",
    category=RuleCategory.DDL,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="单个集合文档数建议不超过 10 亿，超过应考虑分片",
    config={"max_docs": 1_000_000_000},
    message_template="集合文档数超过 {max} 时建议进行分片(sharding)或归档",
))
def check_mongo_col_006(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """集合文档数上限建议。通过 count 操作进行提示。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if success and ("count" in parsed or "estimatedDocumentCount" in parsed):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(max=rule.config["max_docs"]),
        ))
        return
    if re.search(r'\b(countDocuments|estimatedDocumentCount)\s*\(', text):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(max=rule.config["max_docs"]),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-COL-007",
    name="推荐使用 Schema Validator",
    category=RuleCategory.DDL,
    severity=Severity.INFO,
    db_type="mongodb",
    description="建议在创建集合时使用 JSON Schema Validator 进行数据校验",
    message_template="建议为集合添加 JSON Schema Validator 以约束文档结构",
))
def check_mongo_col_007(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """推荐使用 Schema Validator: 检查 create/createCollection 是否包含 validator。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if not success:
        # shell 格式
        if re.search(r'\b(createCollection|create)\s*\(', text):
            if "validator" not in text.lower() and "\"validator\"" not in text:
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code,
                    rule_name=rule.name,
                    level=rule.severity,
                    message=rule.message_template,
                ))
        return
    # JSON 格式: {"create": "coll", ...}
    if "create" in parsed or "createCollection" in parsed:
        if "validator" not in parsed:
            report.findings.append(AuditFindingV2(
                rule_code=rule.code,
                rule_name=rule.name,
                level=rule.severity,
                message=rule.message_template,
            ))


@register_rule(AuditRuleV2(
    code="MONGO-COL-008",
    name="Capped Collection 大小限制",
    category=RuleCategory.DDL,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="Capped Collection 应设置合理的大小限制(默认不超过 1GB)",
    config={"max_capped_size": 1_073_741_824},  # 1GB
    message_template="Capped Collection 大小(size)为 {size} 字节，建议不超过 {max} 字节(1GB)",
))
def check_mongo_col_008(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """Capped Collection 大小限制检查。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if not success:
        return
    create_op = parsed.get("create") or parsed.get("createCollection")
    if not create_op:
        return
    capped = parsed.get("capped", False)
    size = parsed.get("size")
    max_size = rule.config.get("max_capped_size", 1_073_741_824)
    if capped and size and isinstance(size, (int, float)) and size > max_size:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(size=size, max=max_size),
        ))


# ═══════════════════════════════════════════════
# 查询安全规则 (10条) — MONGO-QRY-001 ~ 010
# ═══════════════════════════════════════════════

@register_rule(AuditRuleV2(
    code="MONGO-QRY-001",
    name="find/delete/update 必须带 filter",
    category=RuleCategory.DQL,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="MongoDB 查询/更新/删除操作必须包含 filter 条件，禁止无条件操作",
    message_template="{op} 操作缺少 filter/条件，将操作全集合数据，必须指定查询条件",
))
def check_mongo_qry_001(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """find/delete/update 必须带 filter。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if not success:
        return
    # 检查各类操作
    ops_with_filter = {
        "find": "filter",
        "findOne": "filter",
        "update": "updates",
        "updateOne": "filter",
        "updateMany": "filter",
        "delete": "deletes",
        "deleteOne": "filter",
        "deleteMany": "filter",
    }
    for op, filter_key in ops_with_filter.items():
        if op in parsed:
            # 检查 filter 是否存在且非空
            if filter_key == "filter":
                filter_val = parsed.get(filter_key)
                if filter_val is None or (isinstance(filter_val, dict) and not filter_val):
                    report.findings.append(AuditFindingV2(
                        rule_code=rule.code,
                        rule_name=rule.name,
                        level=rule.severity,
                        message=rule.message_template.format(op=op),
                    ))
                    return
            elif filter_key == "updates":
                updates = parsed.get("updates", [])
                if isinstance(updates, list) and updates:
                    for upd in updates:
                        if isinstance(upd, dict) and not upd.get("q"):
                            report.findings.append(AuditFindingV2(
                                rule_code=rule.code,
                                rule_name=rule.name,
                                level=rule.severity,
                                message=rule.message_template.format(op=op),
                            ))
                            return
            elif filter_key == "deletes":
                deletes = parsed.get("deletes", [])
                if isinstance(deletes, list) and deletes:
                    for d in deletes:
                        if isinstance(d, dict) and not d.get("q"):
                            report.findings.append(AuditFindingV2(
                                rule_code=rule.code,
                                rule_name=rule.name,
                                level=rule.severity,
                                message=rule.message_template.format(op=op),
                            ))
                            return


@register_rule(AuditRuleV2(
    code="MONGO-QRY-002",
    name="推荐指定 projection",
    category=RuleCategory.DQL,
    severity=Severity.INFO,
    db_type="mongodb",
    description="find/findOne 查询推荐指定 projection 限制返回字段，减少网络传输",
    message_template="{op} 查询未指定 projection，建议添加 projection 限制返回字段",
))
def check_mongo_qry_002(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """推荐指定 projection。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if not success:
        return
    query_ops = ("find", "findOne")
    for op in query_ops:
        if op in parsed:
            if "projection" not in parsed:
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code,
                    rule_name=rule.name,
                    level=rule.severity,
                    message=rule.message_template.format(op=op),
                ))
                return


@register_rule(AuditRuleV2(
    code="MONGO-QRY-003",
    name="禁止 $where/$function JavaScript 执行",
    category=RuleCategory.DQL,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="禁止使用 $where 和 $function 运算符执行 JavaScript，存在注入和安全风险",
    message_template="禁止使用 {op} 运算符执行 JavaScript，存在严重的安全和注入风险",
))
def check_mongo_qry_003(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止 $where/$function JavaScript 执行。"""
    text = sql_clean.strip()
    if "\"$where\"" in text or "'$where'" in text:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(op="$where"),
        ))
        return
    if "\"$function\"" in text or "'$function'" in text:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(op="$function"),
        ))
        return
    # 检查函数表达式中的 $where
    if _has_where_operator(text):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(op="$where"),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-QRY-004",
    name="禁止前缀模糊匹配(regex)",
    category=RuleCategory.DQL,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="禁止正则表达式前缀模糊匹配(如 /^abc/)，无法有效使用索引",
    message_template="正则表达式查询 {pattern} 为前缀匹配，无法有效使用索引，建议调整查询策略",
))
def check_mongo_qry_004(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止前缀模糊匹配: 检测 /^ 开头的正则或 $regex 中以 ^ 开头的模式。"""
    text = sql_clean.strip()
    # 检测 $regex 中以 ^ 开头的模式
    regex_patterns = re.findall(r'"\$regex"\s*:\s*"(\^[^"]*)"', text)
    for pat in regex_patterns:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(pattern=pat),
        ))
        return
    # 检测 JS 正则字面量 /^.../
    js_regex = re.findall(r'/\^(.*?)/', text)
    for pat in js_regex:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(pattern=f"^{pat}"),
        ))
        return


@register_rule(AuditRuleV2(
    code="MONGO-QRY-005",
    name="查询返回结果上限",
    category=RuleCategory.DQL,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="查询返回结果建议不超过 10000 条，大数据量应分页",
    config={"max_limit": 10000},
    message_template="查询未设置 limit 或 limit 超过 {max} 条，建议分页查询",
))
def check_mongo_qry_005(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """查询返回结果上限检查。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if not success:
        return
    if "find" not in parsed and "findOne" not in parsed and "aggregate" not in parsed:
        return
    limit = parsed.get("limit")
    max_limit = rule.config.get("max_limit", 10000)
    if limit is None:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(max=max_limit),
        ))
    elif isinstance(limit, (int, float)) and limit > max_limit:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(max=max_limit),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-QRY-006",
    name="explain 不应在生产环境频繁使用",
    category=RuleCategory.DQL,
    severity=Severity.INFO,
    db_type="mongodb",
    description="explain() 会额外收集执行计划信息，不应在生产环境频繁使用",
    message_template="explain() 操作在生产环境频繁使用可能影响性能，建议仅在调试时使用",
))
def check_mongo_qry_006(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """explain() 不应在生产环境频繁使用。"""
    text = sql_clean.strip()
    if "explain" in text.lower() or "\"explain\"" in text:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


@register_rule(AuditRuleV2(
    code="MONGO-QRY-007",
    name="不推荐 snapshot 模式",
    category=RuleCategory.DQL,
    severity=Severity.INFO,
    db_type="mongodb",
    description="snapshot 模式($snapshot 或 readConcern snapshot)性能较差",
    message_template="使用 snapshot 模式可能影响查询性能，如无强一致性需求请避免使用",
))
def check_mongo_qry_007(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """不推荐 snapshot 模式。"""
    text = sql_clean.strip()
    has_snapshot = False
    # JSON 格式
    if "\"snapshot\"" in text and "true" in text.lower():
        has_snapshot = True
    if "snapshot" in text.lower():
        success, parsed = _try_parse_json(text)
        if success:
            snapshot_val = parsed.get("snapshot")
            read_concern = parsed.get("readConcern", {})
            if snapshot_val is True:
                has_snapshot = True
            elif isinstance(read_concern, dict) and read_concern.get("level") == "snapshot":
                has_snapshot = True
    if has_snapshot:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


@register_rule(AuditRuleV2(
    code="MONGO-QRY-008",
    name="不建议使用 allowDiskUse",
    category=RuleCategory.DQL,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="allowDiskUse: true 允许聚合操作使用磁盘，可能表明查询设计不合理",
    message_template="使用 allowDiskUse: true 表明查询可能消耗大量内存，建议优化查询或索引",
))
def check_mongo_qry_008(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """不建议使用 allowDiskUse。"""
    text = sql_clean.strip()
    if "allowDiskUse" not in text:
        return
    success, parsed = _try_parse_json(text)
    if success:
        allow_disk = parsed.get("allowDiskUse")
        if allow_disk is True:
            report.findings.append(AuditFindingV2(
                rule_code=rule.code,
                rule_name=rule.name,
                level=rule.severity,
                message=rule.message_template,
            ))
            return
    # 正则匹配
    if re.search(r'"allowDiskUse"\s*:\s*true', text, re.IGNORECASE):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


@register_rule(AuditRuleV2(
    code="MONGO-QRY-009",
    name="text 搜索使用限制",
    category=RuleCategory.DQL,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="MongoDB 的 $text 搜索功能有限，建议使用 Elasticsearch 等专业搜索引擎",
    message_template="使用 $text 全文搜索功能有限，建议考虑 Elasticsearch 等专业搜索方案",
))
def check_mongo_qry_009(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """text 搜索使用限制。"""
    text = sql_clean.strip()
    if "\"$text\"" in text or "'$text'" in text:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


@register_rule(AuditRuleV2(
    code="MONGO-QRY-010",
    name="filter 必须是有效的 BSON 格式",
    category=RuleCategory.DQL,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="MongoDB filter 必须是有效的 JSON/BSON 对象格式",
    message_template="查询条件 filter 不是有效的 JSON/BSON 对象格式，请修正",
))
def check_mongo_qry_010(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """filter 必须是有效的 BSON 格式。"""
    text = sql_clean.strip()
    if not text.startswith("{"):
        return
    # 尝试解析 JSON
    success, parsed = _try_parse_json(text)
    if not success:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))
        return
    # 检查 filter 字段是否为对象
    for filter_key in ("filter", "q"):
        if filter_key in parsed:
            filter_val = parsed[filter_key]
            if filter_val is not None and not isinstance(filter_val, dict):
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code,
                    rule_name=rule.name,
                    level=rule.severity,
                    message=rule.message_template,
                ))
                return


# ═══════════════════════════════════════════════
# 聚合管道规则 (6条) — MONGO-AGG-001 ~ 006
# ═══════════════════════════════════════════════

@register_rule(AuditRuleV2(
    code="MONGO-AGG-001",
    name="禁止 $out/$merge 写操作",
    category=RuleCategory.DQL,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="聚合管道中禁止使用 $out 和 $merge 阶段，可能覆盖现有集合数据",
    message_template="禁止在聚合管道中使用 {stage} 阶段，可能覆盖或修改现有集合数据",
))
def check_mongo_agg_001(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止 $out/$merge 写操作。"""
    text = sql_clean.strip()
    # 检查 JSON 中的 pipeline
    success, parsed = _try_parse_json(text)
    if success and "pipeline" in parsed and isinstance(parsed["pipeline"], list):
        for stage in parsed["pipeline"]:
            if isinstance(stage, dict):
                if "$out" in stage:
                    out_val = stage["$out"]
                    report.findings.append(AuditFindingV2(
                        rule_code=rule.code,
                        rule_name=rule.name,
                        level=rule.severity,
                        message=rule.message_template.format(stage=f"$out: {out_val}"),
                    ))
                    return
                if "$merge" in stage:
                    report.findings.append(AuditFindingV2(
                        rule_code=rule.code,
                        rule_name=rule.name,
                        level=rule.severity,
                        message=rule.message_template.format(stage="$merge"),
                    ))
                    return
    # 文本检查
    if "\"$out\"" in text or "'$out'" in text:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(stage="$out"),
        ))
        return
    if "\"$merge\"" in text or "'$merge'" in text:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(stage="$merge"),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-AGG-002",
    name="聚合管道阶段数上限",
    category=RuleCategory.DQL,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="聚合管道阶段数建议不超过 20 个，过多阶段影响性能",
    config={"max_stages": 20},
    message_template="聚合管道包含 {count} 个阶段，超过建议上限 {max} 个",
))
def check_mongo_agg_002(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """聚合管道阶段数上限检查。"""
    text = sql_clean.strip()
    stage_count = 0
    success, parsed = _try_parse_json(text)
    if success and "pipeline" in parsed and isinstance(parsed["pipeline"], list):
        stage_count = len(parsed["pipeline"])
    else:
        # 粗略统计 $ 开头的 stage 数量
        stage_count = len(re.findall(r'\{\\s*"\$(\w+)"', text))
    max_stages = rule.config.get("max_stages", 20)
    if stage_count > max_stages:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(count=stage_count, max=max_stages),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-AGG-003",
    name="$unwind 大数据量警告",
    category=RuleCategory.DQL,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="$unwind 展开数组可能导致文档数量急剧膨胀，影响性能",
    message_template="使用 $unwind 阶段可能急剧增加文档数量，请确认数组字段的平均元素数量",
))
def check_mongo_agg_003(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """$unwind 大数据量警告。"""
    text = sql_clean.strip()
    if "\"$unwind\"" in text or "'$unwind'" in text:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


@register_rule(AuditRuleV2(
    code="MONGO-AGG-004",
    name="$group 操作索引建议",
    category=RuleCategory.DQL,
    severity=Severity.INFO,
    db_type="mongodb",
    description="$group 操作的 _id 字段建议创建索引以提升性能",
    message_template="$group 操作的 _id 字段建议创建索引，可显著提升聚合性能",
))
def check_mongo_agg_004(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """$group 操作索引建议。"""
    text = sql_clean.strip()
    if "\"$group\"" in text or "'$group'" in text:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


@register_rule(AuditRuleV2(
    code="MONGO-AGG-005",
    name="聚合内存限制(100MB)",
    category=RuleCategory.DQL,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="单个聚合阶段内存限制为 100MB，超出需使用 allowDiskUse",
    config={"memory_limit_mb": 100},
    message_template="聚合操作内存限制为 {limit}MB，大型聚合可能触发内存限制错误",
))
def check_mongo_agg_005(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """聚合内存限制提示。"""
    text = sql_clean.strip()
    if "aggregate" in text.lower() or "\"pipeline\"" in text:
        limit_mb = rule.config.get("memory_limit_mb", 100)
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(limit=limit_mb),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-AGG-006",
    name="$lookup 嵌套层数限制",
    category=RuleCategory.DQL,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="$lookup 嵌套层数建议不超过 3 层，过多嵌套影响性能",
    config={"max_lookup_nesting": 3},
    message_template="$lookup 嵌套层数可能超过 {max} 层，建议优化查询结构",
))
def check_mongo_agg_006(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """$lookup 嵌套层数限制。"""
    text = sql_clean.strip()
    lookup_count = text.count('"$lookup"') + text.count("'$lookup'")
    max_nesting = rule.config.get("max_lookup_nesting", 3)
    if lookup_count > max_nesting:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(max=max_nesting),
        ))


# ═══════════════════════════════════════════════
# 写入安全规则 (8条) — MONGO-WR-001 ~ 008
# ═══════════════════════════════════════════════

@register_rule(AuditRuleV2(
    code="MONGO-WR-001",
    name="批量写入上限(1000条)",
    category=RuleCategory.DML,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="insertMany/bulkWrite 单次文档数建议不超过 1000 条",
    config={"max_bulk_write": 1000},
    message_template="批量写入文档数为 {count}，超过建议上限 {max} 条，请分批写入",
))
def check_mongo_wr_001(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """批量写入上限检查。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    doc_count = 0
    if success:
        # insertMany: {"documents": [...]}
        if "insertMany" in parsed or "insert" in parsed:
            docs = parsed.get("documents", [])
            if isinstance(docs, list):
                doc_count = len(docs)
        # bulkWrite: {"operations": [...]}
        if "operations" in parsed and isinstance(parsed["operations"], list):
            doc_count = len(parsed["operations"])
        # updateMany: {"updates": [...]}
        if "updates" in parsed and isinstance(parsed["updates"], list):
            doc_count = len(parsed["updates"])
    max_docs = rule.config.get("max_bulk_write", 1000)
    if doc_count > max_docs:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(count=doc_count, max=max_docs),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-WR-002",
    name="updateMany 必须带条件 filter",
    category=RuleCategory.DML,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="updateMany 操作必须包含 filter 条件，禁止无条件更新",
    message_template="updateMany 操作缺少有效的 filter 条件，将更新全集合文档",
))
def check_mongo_wr_002(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """updateMany 必须带条件 filter。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if not success:
        return
    if "updateMany" in parsed:
        filter_val = parsed.get("filter")
        if filter_val is None or (isinstance(filter_val, dict) and not filter_val):
            report.findings.append(AuditFindingV2(
                rule_code=rule.code,
                rule_name=rule.name,
                level=rule.severity,
                message=rule.message_template,
            ))
    # 检查 bulk updates
    if "updates" in parsed and isinstance(parsed.get("updates"), list):
        for upd in parsed["updates"]:
            if isinstance(upd, dict):
                q = upd.get("q")
                u = upd.get("u")
                multi = upd.get("multi", False)
                if multi and (q is None or (isinstance(q, dict) and not q)):
                    report.findings.append(AuditFindingV2(
                        rule_code=rule.code,
                        rule_name=rule.name,
                        level=rule.severity,
                        message=rule.message_template,
                    ))
                    return


@register_rule(AuditRuleV2(
    code="MONGO-WR-003",
    name="deleteMany 必须带条件 filter",
    category=RuleCategory.DML,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="deleteMany 操作必须包含 filter 条件，禁止无条件删除",
    message_template="deleteMany 操作缺少有效的 filter 条件，将删除全集合文档",
))
def check_mongo_wr_003(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """deleteMany 必须带条件 filter。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if not success:
        return
    if "deleteMany" in parsed:
        filter_val = parsed.get("filter")
        if filter_val is None or (isinstance(filter_val, dict) and not filter_val):
            report.findings.append(AuditFindingV2(
                rule_code=rule.code,
                rule_name=rule.name,
                level=rule.severity,
                message=rule.message_template,
            ))
    # 检查 bulk deletes
    if "deletes" in parsed and isinstance(parsed.get("deletes"), list):
        for d in parsed["deletes"]:
            if isinstance(d, dict):
                q = d.get("q")
                limit = d.get("limit", 0)
                if limit == 0 and (q is None or (isinstance(q, dict) and not q)):
                    report.findings.append(AuditFindingV2(
                        rule_code=rule.code,
                        rule_name=rule.name,
                        level=rule.severity,
                        message=rule.message_template,
                    ))
                    return


@register_rule(AuditRuleV2(
    code="MONGO-WR-004",
    name="禁止 unset 关键字段",
    category=RuleCategory.DML,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="禁止 $unset 关键字段(如 _id, create_time, update_time 等)",
    config={"critical_fields": ["_id", "create_time", "update_time", "created_at", "updated_at"]},
    message_template="禁止使用 $unset 删除关键字段: {fields}，可能导致数据完整性问题",
))
def check_mongo_wr_004(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止 unset 关键字段。"""
    text = sql_clean.strip()
    critical = rule.config.get("critical_fields", ["_id"])
    if "\"$unset\"" not in text and "'$unset'" not in text:
        return
    success, parsed = _try_parse_json(text)
    found_fields: List[str] = []
    if success:
        # 查找 updates 中的 $unset
        for update_dict in _find_all_values(parsed, "u"):
            if isinstance(update_dict, dict) and "$unset" in update_dict:
                unset_fields = update_dict["$unset"]
                if isinstance(unset_fields, dict):
                    for f in unset_fields:
                        if f in critical:
                            found_fields.append(f)
                elif isinstance(unset_fields, list):
                    for f in unset_fields:
                        if f in critical:
                            found_fields.append(f)
        # 也检查顶层 update 的 $unset
        if isinstance(parsed, dict):
            for k, v in parsed.items():
                if isinstance(v, dict) and "$unset" in v:
                    unset_val = v["$unset"]
                    if isinstance(unset_val, dict):
                        for f in unset_val:
                            if f in critical:
                                found_fields.append(f)
    else:
        # 正则提取
        unset_match = re.findall(r'"\$unset"\s*:\s*\{([^}]+)\}', text)
        for m in unset_match:
            for f in critical:
                if f'"{f}"' in m or f"'{f}'" in m:
                    found_fields.append(f)
    if found_fields:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(fields=", ".join(set(found_fields))),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-WR-005",
    name="禁止重命名集合",
    category=RuleCategory.DML,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="禁止执行 renameCollection 操作，可能影响应用连接",
    message_template="禁止执行 renameCollection 操作，重命名集合可能导致应用连接失败",
))
def check_mongo_wr_005(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止重命名集合。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if success:
        if "renameCollection" in parsed or "collMod" in parsed:
            report.findings.append(AuditFindingV2(
                rule_code=rule.code,
                rule_name=rule.name,
                level=rule.severity,
                message=rule.message_template,
            ))
            return
    if re.search(r'\b(renameCollection|collMod)\b', text, re.IGNORECASE):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


@register_rule(AuditRuleV2(
    code="MONGO-WR-006",
    name="禁止绕过文档验证(bypassDocumentValidation)",
    category=RuleCategory.DML,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="禁止设置 bypassDocumentValidation: true 绕过文档验证",
    message_template="禁止设置 bypassDocumentValidation: true 绕过 Schema 验证",
))
def check_mongo_wr_006(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止绕过文档验证。"""
    text = sql_clean.strip()
    if "bypassDocumentValidation" not in text:
        return
    success, parsed = _try_parse_json(text)
    if success:
        bypass = parsed.get("bypassDocumentValidation")
        if bypass is True:
            report.findings.append(AuditFindingV2(
                rule_code=rule.code,
                rule_name=rule.name,
                level=rule.severity,
                message=rule.message_template,
            ))
            return
    # 正则匹配
    if re.search(r'"bypassDocumentValidation"\s*:\s*true', text, re.IGNORECASE):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


@register_rule(AuditRuleV2(
    code="MONGO-WR-007",
    name="单文档大小限制(16MB警告阈值10MB)",
    category=RuleCategory.DML,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="MongoDB 单文档最大 16MB，建议不超过 10MB 作为警告阈值",
    config={"warn_threshold_bytes": 10_485_760, "max_size_bytes": 16_777_216},  # 10MB, 16MB
    message_template="文档大小约为 {size_mb}MB，超过警告阈值 {warn_mb}MB，建议拆分或精简文档",
))
def check_mongo_wr_007(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """单文档大小限制检查。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if not success:
        return
    # 检查是否为写入操作
    write_ops = ("insert", "insertOne", "insertMany", "update", "updateOne", "updateMany")
    is_write = any(op in parsed for op in write_ops)
    if not is_write:
        return
    # 估算文档大小
    doc_size = _get_json_string_length(text)
    warn_threshold = rule.config.get("warn_threshold_bytes", 10_485_760)
    if doc_size > warn_threshold:
        size_mb = round(doc_size / (1024 * 1024), 2)
        warn_mb = round(warn_threshold / (1024 * 1024), 2)
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(size_mb=size_mb, warn_mb=warn_mb),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-WR-008",
    name="禁止嵌套深度超过5层",
    category=RuleCategory.DML,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="文档嵌套深度建议不超过 5 层，过深影响查询和更新性能",
    config={"max_nesting_depth": 5},
    message_template="文档嵌套深度为 {depth} 层，超过建议上限 {max} 层",
))
def check_mongo_wr_008(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止嵌套深度超过5层。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if not success:
        return
    # 检查 documents 或 document 字段的嵌套深度
    docs = parsed.get("documents") or parsed.get("document")
    if docs is None:
        # 检查 updates 中的 $set/$setOnInsert 等
        updates = parsed.get("updates", [])
        if isinstance(updates, list) and updates:
            docs = updates[0].get("u") if isinstance(updates[0], dict) else None
    if not docs:
        return
    target_doc = docs[0] if isinstance(docs, list) else docs
    if not isinstance(target_doc, dict):
        return
    depth = _get_nesting_depth(target_doc)
    max_depth = rule.config.get("max_nesting_depth", 5)
    if depth > max_depth:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(depth=depth, max=max_depth),
        ))


# ═══════════════════════════════════════════════
# 安全规则 (5条) — MONGO-SEC-001 ~ 005
# ═══════════════════════════════════════════════

@register_rule(AuditRuleV2(
    code="MONGO-SEC-001",
    name="禁止 db.eval() 执行 JavaScript",
    category=RuleCategory.SECURITY,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="禁止 db.eval() 执行任意 JavaScript 代码，存在严重安全风险",
    message_template="禁止 db.eval() 执行 JavaScript 代码，存在严重的安全和注入风险",
))
def check_mongo_sec_001(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止 db.eval() 执行 JavaScript。"""
    text = sql_clean.strip()
    if re.search(r'\bdb\.eval\s*\(', text):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


@register_rule(AuditRuleV2(
    code="MONGO-SEC-002",
    name="禁止操作系统集合",
    category=RuleCategory.SECURITY,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="禁止查询/操作系统集合(以 system. 开头)，可能导致安全问题",
    message_template="禁止操作系统集合: '{coll}'，系统集合操作可能导致安全问题",
))
def check_mongo_sec_002(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止操作系统集合。"""
    text = sql_clean.strip()
    coll = _extract_collection_name(text)
    if coll and coll.startswith("system."):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(coll=coll),
        ))
        return
    # 也检查 JSON 中的 find 等操作
    success, parsed = _try_parse_json(text)
    if success:
        for op_key in ("find", "findOne", "aggregate", "insert", "update", "delete", "count"):
            val = parsed.get(op_key)
            if isinstance(val, str) and val.startswith("system."):
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code,
                    rule_name=rule.name,
                    level=rule.severity,
                    message=rule.message_template.format(coll=val),
                ))
                return


@register_rule(AuditRuleV2(
    code="MONGO-SEC-003",
    name="禁止 MapReduce(使用聚合管道)",
    category=RuleCategory.SECURITY,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="禁止使用 mapReduce 命令，map 和 reduce 函数可执行任意 JavaScript",
    message_template="禁止 mapReduce 操作，请使用聚合管道(Aggregation Pipeline)替代",
))
def check_mongo_sec_003(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止 MapReduce。"""
    text = sql_clean.strip()
    if re.search(r'\bmapReduce\b|\bmapReduce\s*\(', text) or "\"mapReduce\"" in text:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


@register_rule(AuditRuleV2(
    code="MONGO-SEC-004",
    name="防止 $cmd 注入攻击",
    category=RuleCategory.SECURITY,
    severity=Severity.ERROR,
    db_type="mongodb",
    description="禁止直接操作 $cmd 集合，可能存在注入攻击风险",
    message_template="禁止直接操作 $cmd 集合，存在注入攻击风险",
))
def check_mongo_sec_004(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """防止 $cmd 注入攻击。"""
    text = sql_clean.strip()
    if "$cmd" in text:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


@register_rule(AuditRuleV2(
    code="MONGO-SEC-005",
    name="敏感操作必须记录审计日志",
    category=RuleCategory.SECURITY,
    severity=Severity.INFO,
    db_type="mongodb",
    description="删除、更新、聚合写操作等敏感操作应记录审计日志",
    message_template="敏感操作({ops})应确保已记录审计日志，便于安全追踪和合规审计",
))
def check_mongo_sec_005(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """敏感操作审计日志提醒。"""
    text = sql_clean.strip()
    sensitive_ops = []
    success, parsed = _try_parse_json(text)
    op_names = {
        "delete": "delete", "deleteOne": "delete", "deleteMany": "delete",
        "update": "update", "updateOne": "update", "updateMany": "update",
        "drop": "drop", "dropDatabase": "dropDatabase", "dropIndexes": "dropIndexes",
        "renameCollection": "rename", "createUser": "userManagement",
        "dropUser": "userManagement", "grantRoles": "userManagement",
    }
    if success:
        for op, label in op_names.items():
            if op in parsed:
                sensitive_ops.append(label)
    else:
        text_lower = text.lower()
        for op, label in op_names.items():
            if op.lower() in text_lower:
                sensitive_ops.append(label)
    if sensitive_ops:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(ops=", ".join(set(sensitive_ops))),
        ))


# ═══════════════════════════════════════════════
# 性能规则 (5条) — MONGO-PERF-001 ~ 005
# ═══════════════════════════════════════════════

@register_rule(AuditRuleV2(
    code="MONGO-PERF-001",
    name="缺少索引查询警告",
    category=RuleCategory.PERFORMANCE,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="查询条件字段应创建索引，缺少索引会导致全集合扫描(COLLSCAN)",
    message_template="查询条件字段 '{fields}' 缺少索引可能导致全集合扫描，建议创建索引",
))
def check_mongo_perf_001(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """缺少索引查询警告: 检查查询条件中常见字段是否可能缺少索引。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if not success:
        return
    if "find" not in parsed and "findOne" not in parsed:
        return
    filter_val = parsed.get("filter")
    if not isinstance(filter_val, dict) or not filter_val:
        return
    # 检查常见的可能需要索引的字段（排除 _id 已有默认索引）
    query_fields = [k for k in filter_val.keys() if not k.startswith("$") and k != "_id"]
    if query_fields:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template.format(fields=", ".join(query_fields)),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-PERF-002",
    name="推荐覆盖索引查询",
    category=RuleCategory.PERFORMANCE,
    severity=Severity.INFO,
    db_type="mongodb",
    description="推荐 projection 仅包含索引字段，实现覆盖索引查询(Covered Query)",
    message_template="查询条件 '{fields}' 建议 projection 仅包含索引字段以实现覆盖索引查询",
))
def check_mongo_perf_002(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """推荐覆盖索引查询。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if not success:
        return
    if "find" not in parsed and "findOne" not in parsed:
        return
    filter_val = parsed.get("filter")
    projection = parsed.get("projection")
    if not isinstance(filter_val, dict) or not filter_val:
        return
    # 如果有 projection 且设置了 _id: 0，可能是尝试覆盖索引
    if isinstance(projection, dict) and projection.get("_id") == 0:
        query_fields = [k for k in filter_val.keys() if not k.startswith("$")]
        if query_fields:
            report.findings.append(AuditFindingV2(
                rule_code=rule.code,
                rule_name=rule.name,
                level=rule.severity,
                message=rule.message_template.format(fields=", ".join(query_fields)),
            ))


@register_rule(AuditRuleV2(
    code="MONGO-PERF-003",
    name="全集合扫描警告",
    category=RuleCategory.PERFORMANCE,
    severity=Severity.WARNING,
    db_type="mongodb",
    description="缺少 filter 条件或 filter 为空对象 {} 将导致全集合扫描",
    message_template="查询缺少有效过滤条件，将导致全集合扫描(COLLSCAN)，严重影响性能",
))
def check_mongo_perf_003(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """全集合扫描警告。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if not success:
        return
    if "find" not in parsed and "findOne" not in parsed:
        return
    filter_val = parsed.get("filter")
    if filter_val is None:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))
    elif isinstance(filter_val, dict) and not filter_val:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


@register_rule(AuditRuleV2(
    code="MONGO-PERF-004",
    name="复合索引字段顺序建议",
    category=RuleCategory.PERFORMANCE,
    severity=RuleCategory.PERFORMANCE,
    db_type="mongodb",
    description="复合索引字段顺序应遵循 ESR 规则(Equality, Sort, Range)",
    message_template="复合查询条件 '{fields}' 建议遵循 ESR 规则设计索引: 等值 > 排序 > 范围",
))
def check_mongo_perf_004(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """复合索引字段顺序建议: ESR 规则。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if not success:
        return
    if "find" not in parsed and "findOne" not in parsed:
        return
    filter_val = parsed.get("filter")
    if not isinstance(filter_val, dict) or not filter_val:
        return
    # 统计查询条件中的字段数量
    query_fields = [k for k in filter_val.keys() if not k.startswith("$")]
    # 如果有多于 1 个查询字段，提示复合索引 ESR 规则
    if len(query_fields) > 1:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=Severity.INFO,
            message=rule.message_template.format(fields=", ".join(query_fields)),
        ))


@register_rule(AuditRuleV2(
    code="MONGO-PERF-005",
    name="分片键选择建议",
    category=RuleCategory.PERFORMANCE,
    severity=Severity.INFO,
    db_type="mongodb",
    description="分片集合应选择高基数、均匀分布的字段作为分片键",
    message_template="分片集合建议: 选择高基数均匀分布字段作为分片键，避免单调递增字段和热点",
))
def check_mongo_perf_005(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """分片键选择建议: 检查 shardCollection 操作。"""
    text = sql_clean.strip()
    success, parsed = _try_parse_json(text)
    if success and "shardCollection" in parsed:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))
        return
    if re.search(r'\bshardCollection\b', text):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code,
            rule_name=rule.name,
            level=rule.severity,
            message=rule.message_template,
        ))


# ──────────────────────────────────────────────
# 规则统计信息
# ──────────────────────────────────────────────


def get_mongodb_rule_count() -> Dict[str, int]:
    """获取 MongoDB 规则分类统计。

    Returns:
        {"COL": 8, "QRY": 10, "AGG": 6, "WR": 8, "SEC": 5, "PERF": 5, "TOTAL": 42}
    """
    return {
        "COL": 8,
        "QRY": 10,
        "AGG": 6,
        "WR": 8,
        "SEC": 5,
        "PERF": 5,
        "TOTAL": 42,
    }