"""PostgreSQL 特有审核规则引擎。

包含 45+ 条 PostgreSQL 特有审核规则，覆盖以下维度:
- DDL (18条): CREATE TABLE / ALTER TABLE 特有语法规范
- DML (8条): INSERT/UPDATE/DELETE 操作规范
- DQL (6条): SELECT 查询规范
- SECURITY (6条): PostgreSQL 安全风险检测
- PERFORMANCE (7条): PostgreSQL 性能优化建议

规则编码规范:
- DDL: PG-DDL-001 ~ PG-DDL-018
- DML: PG-DML-001 ~ PG-DML-008
- DQL: PG-DQL-001 ~ PG-DQL-006
- SEC: PG-SEC-001 ~ PG-SEC-006
- PERF: PG-PERF-001 ~ PG-PERF-007

依赖:
    from app.services.audit_engine_v2 import (
        AuditRuleV2, AuditFindingV2, AuditReportV2,
        Severity, RuleCategory, register_rule,
        _extract_table_name, _extract_column_definitions,
        _extract_index_definitions,
    )
"""

from __future__ import annotations

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

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


# ═══════════════════════════════════════════════
# PostgreSQL 规则专用辅助函数
# ═══════════════════════════════════════════════


def _extract_pg_table_name(sql_upper: str, sql_clean: str) -> Optional[str]:
    """从 PostgreSQL CREATE/ALTER TABLE 语句中提取表名（支持 schema.table 形式）。"""
    m = re.search(
        r"\bCREATE\s+(?:TEMPORARY\s+|TEMP\s+)?TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?(?:\"?\w+\"?\.)?(\"?\w+\"?)",
        sql_clean, re.IGNORECASE,
    )
    if m:
        return m.group(1).strip('"')
    m = re.search(
        r"\bALTER\s+TABLE\s+(?:ONLY\s+)?(?:\"?\w+\"?\.)?(\"?\w+\"?)",
        sql_clean, re.IGNORECASE,
    )
    if m:
        return m.group(1).strip('"')
    return None


def _extract_pg_column_defs(sql_upper: str) -> List[Dict[str, str]]:
    """从 PostgreSQL CREATE TABLE (...) 中提取列定义。"""
    paren_match = re.search(r"\((.*)\)\s*(?:INHERITS|TABLESPACE|PARTITION|WITH|;|$)", sql_upper, re.DOTALL)
    if not paren_match:
        return []
    inner = paren_match.group(1)
    parts = []
    depth = 0
    current = ""
    for char in inner:
        if char == "(":
            depth += 1
            current += char
        elif char == ")":
            depth -= 1
            current += char
        elif char == "," and depth == 0:
            parts.append(current.strip())
            current = ""
        else:
            current += char
    if current.strip():
        parts.append(current.strip())
    columns: List[Dict[str, str]] = []
    for part in parts:
        if re.match(r"^(CONSTRAINT|PRIMARY\s+KEY|UNIQUE|CHECK|EXCLUDE|FOREIGN\s+KEY|REFERENCES|LIKE)\b", part, re.IGNORECASE):
            continue
        col_match = re.match(r'"?(\w+)"?\s+(\w+(?:\s*\[\s*\])?|\w+\s+\w+)', part, re.IGNORECASE)
        if col_match:
            columns.append({"name": col_match.group(1), "type": col_match.group(2).upper(), "full": part})
    return columns


def _has_pg_serial_type(col_type: str) -> bool:
    """判断是否为 SERIAL 类型。"""
    return any(st in col_type for st in ("SERIAL", "BIGSERIAL", "SMALLSERIAL"))


# ═══════════════════════════════════════════════
# DDL 规则 (PG-DDL-001 ~ PG-DDL-018, 共18条)
# ═══════════════════════════════════════════════


# ── PG-DDL-001 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-001", name="PGDDLEnablePrimaryKey",
    category=RuleCategory.DDL, severity=Severity.ERROR, db_type="postgres",
    description="CREATE TABLE 语句必须包含 PRIMARY KEY 定义",
    config={"allow_no_pk": False},
    message_template="建表语句缺少主键定义，请添加 PRIMARY KEY（推荐使用 GENERATED ALWAYS AS IDENTITY）",
))
def check_pg_ddl_001(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """强制必须拥有主键。"""
    if not re.search(r"\bCREATE\s+TABLE\b", sql_upper, re.IGNORECASE):
        return
    if "PRIMARY KEY" not in sql_upper:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
        ))


# ── PG-DDL-002 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-002", name="PGDDLCheckTableComment",
    category=RuleCategory.DDL, severity=Severity.WARNING, db_type="postgres",
    description="建表语句必须包含 COMMENT ON 表注释",
    config={"require_column_comment": True},
    message_template="表 '{table}' 缺少注释，请添加 COMMENT ON TABLE {table} IS '...'",
))
def check_pg_ddl_002(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """强制表/列必须拥有注释(COMMENT ON)。"""
    table_name = _extract_pg_table_name(sql_upper, sql_clean)
    if not table_name:
        return
    has_comment = re.search(r"\bCOMMENT\s+ON\s+(?:TABLE|COLUMN)\b", sql_upper, re.IGNORECASE)
    if not has_comment:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(table=table_name),
        ))


# ── PG-DDL-003 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-003", name="PGDDLCheckColumnNullable",
    category=RuleCategory.DDL, severity=Severity.WARNING, db_type="postgres",
    description="非timestamp类型字段必须声明 NOT NULL",
    config={"nullable_exceptions": ["TIMESTAMP", "TIMESTAMPTZ", "DATE", "TIME", "JSON", "JSONB", "TEXT"]},
    message_template="字段 '{column}' 类型为 {col_type}，未声明 NOT NULL，建议添加约束",
))
def check_pg_ddl_003(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """强制非timestamp类型字段NOT NULL。"""
    if not re.search(r"\bCREATE\s+TABLE\b", sql_upper, re.IGNORECASE):
        return
    columns = _extract_pg_column_defs(sql_upper)
    exceptions = rule.config.get("nullable_exceptions", [])
    for col in columns:
        col_type = col["type"]
        if any(exc in col_type for exc in exceptions):
            continue
        if "NOT NULL" not in col["full"] and "NULL" not in col["full"]:
            if "DEFAULT" in col["full"]:
                continue
            if _has_pg_serial_type(col_type):
                continue
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                message=rule.message_template.format(column=col["name"], col_type=col_type),
            ))


# ── PG-DDL-004 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-004", name="PGDDLCheckColumnDefault",
    category=RuleCategory.DDL, severity=Severity.WARNING, db_type="postgres",
    description="字段应当拥有默认值（除主键和 GENERATED 列外）",
    config={"require_default": True},
    message_template="字段 '{column}' 缺少 DEFAULT 值，建议添加",
))
def check_pg_ddl_004(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """字段必须拥有默认值。"""
    if not re.search(r"\bCREATE\s+TABLE\b", sql_upper, re.IGNORECASE):
        return
    for col in _extract_pg_column_defs(sql_upper):
        if "PRIMARY KEY" in col["full"] or "GENERATED" in col["full"] or _has_pg_serial_type(col["type"]):
            continue
        if "DEFAULT" not in col["full"]:
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                message=rule.message_template.format(column=col["name"]),
            ))


# ── PG-DDL-005 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-005", name="PGDDLCheckFloatDouble",
    category=RuleCategory.DDL, severity=Severity.WARNING, db_type="postgres",
    description="浮点类型精度丢失，建议使用 NUMERIC/DECIMAL",
    message_template="字段 '{column}' 使用了 {col_type} 类型，建议使用 NUMERIC(p,s) 以确保精度",
))
def check_pg_ddl_005(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """NUMERIC vs FLOAT 精度建议。"""
    if not re.search(r"\bCREATE\s+TABLE\b|\bALTER\s+TABLE\b", sql_upper, re.IGNORECASE):
        return
    float_types = ("FLOAT", "REAL", "DOUBLE", "DOUBLE PRECISION")
    for col in _extract_pg_column_defs(sql_upper):
        for ft in float_types:
            if ft in col["type"]:
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                    message=rule.message_template.format(column=col["name"], col_type=col["type"]),
                ))
                break


# ── PG-DDL-006 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-006", name="PGDDLEnableSerial",
    category=RuleCategory.DDL, severity=Severity.INFO, db_type="postgres",
    description="推荐使用 GENERATED ALWAYS AS IDENTITY 替代 SERIAL",
    config={"recommend_identity": True, "generated_type": "ALWAYS"},
    message_template="建议使用 GENERATED {generated_type} AS IDENTITY 替代 SERIAL/BIGSERIAL，更符合 SQL 标准",
))
def check_pg_ddl_006(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """推荐自增列使用 GENERATED ALWAYS AS IDENTITY。"""
    if not re.search(r"\bCREATE\s+TABLE\b", sql_upper, re.IGNORECASE):
        return
    for col in _extract_pg_column_defs(sql_upper):
        if _has_pg_serial_type(col["type"]):
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                message=rule.message_template.format(generated_type=rule.config.get("generated_type", "ALWAYS")),
            ))
            break


# ── PG-DDL-007 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-007", name="PGDDLPrimaryKeyMust",
    category=RuleCategory.DDL, severity=Severity.INFO, db_type="postgres",
    description="主键字段名应为 id（支持配置）",
    config={"pk_name": "id"},
    message_template="主键名 '{pk_name}' 建议改为 '{suggested}'",
))
def check_pg_ddl_007(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """主键名称为 id。"""
    if not re.search(r"\bCREATE\s+TABLE\b", sql_upper, re.IGNORECASE):
        return
    suggested = rule.config.get("pk_name", "id")
    pk_match = re.search(r"PRIMARY\s+KEY\s*\(([^)]+)\)", sql_upper, re.IGNORECASE)
    if pk_match:
        for col in [c.strip().strip('"') for c in pk_match.group(1).split(",")]:
            if col.lower() != suggested.lower():
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                    message=rule.message_template.format(pk_name=col, suggested=suggested),
                ))
    for col in _extract_pg_column_defs(sql_upper):
        if "PRIMARY KEY" in col["full"] and col["name"].lower() != suggested.lower():
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                message=rule.message_template.format(pk_name=col["name"], suggested=suggested),
            ))


# ── PG-DDL-008 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-008", name="PGDDLEnableDropTable",
    category=RuleCategory.DDL, severity=Severity.ERROR, db_type="postgres",
    description="生产环境默认禁止执行 DROP TABLE 操作",
    config={"env_production_only": True, "allow_drop": False},
    message_template="禁止执行 DROP TABLE 操作，数据将永久丢失。如需删除请先备份",
))
def check_pg_ddl_008(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """允许删除表(生产环境默认禁止)。"""
    if not re.search(r"\bDROP\s+TABLE\b", sql_upper, re.IGNORECASE):
        return
    if not rule.config.get("allow_drop", False):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
        ))


# ── PG-DDL-009 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-009", name="PGDDLEnableDropDatabase",
    category=RuleCategory.DDL, severity=Severity.ERROR, db_type="postgres",
    description="禁止执行 DROP DATABASE 操作",
    config={"allow_drop_db": False},
    message_template="禁止执行 DROP DATABASE 操作，将导致整个数据库被删除",
))
def check_pg_ddl_009(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """允许删除库（默认禁止）。"""
    if not re.search(r"\bDROP\s+DATABASE\b", sql_upper, re.IGNORECASE):
        return
    if not rule.config.get("allow_drop_db", False):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
        ))


# ── PG-DDL-010 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-010", name="PGDDLAllowPRINotInt",
    category=RuleCategory.DDL, severity=Severity.INFO, db_type="postgres",
    description="检查主键类型是否为 integer/bigint（允许 UUID 等）",
    config={"allowed_pk_types": ["INT", "BIGINT", "INTEGER", "SERIAL", "BIGSERIAL", "UUID"]},
    message_template="主键类型 '{pk_type}' 不在推荐列表中，建议使用 BIGINT/SERIAL/UUID",
))
def check_pg_ddl_010(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """允许主键类型为非integer/bigint（给出建议）。"""
    if not re.search(r"\bCREATE\s+TABLE\b", sql_upper, re.IGNORECASE):
        return
    allowed = rule.config.get("allowed_pk_types", [])
    for col in _extract_pg_column_defs(sql_upper):
        if "PRIMARY KEY" in col["full"]:
            if not any(a.upper() in col["type"] for a in allowed):
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                    message=rule.message_template.format(pk_type=col["type"]),
                ))


# ── PG-DDL-011 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-011", name="PGDDLMultiToCommit",
    category=RuleCategory.DDL, severity=Severity.INFO, db_type="postgres",
    description="单工单包含多条 DDL 语句",
    config={"max_ddl_per_ticket": 5, "allow_multi": True},
    message_template="检测到单工单包含 {count} 条 DDL 语句，建议不超过 {max} 条",
))
def check_pg_ddl_011(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """允许单工单多条DDL（但建议不超过限制）。"""
    ddl_patterns = [
        r"\bCREATE\s+(?:TABLE|INDEX|SEQUENCE|VIEW|MATERIALIZED\s+VIEW|TRIGGER|FUNCTION)\b",
        r"\bALTER\s+TABLE\b",
        r"\bDROP\s+(?:TABLE|INDEX|SEQUENCE|VIEW)\b",
    ]
    count = sum(len(re.findall(pat, sql_upper, re.IGNORECASE)) for pat in ddl_patterns)
    max_ddl = rule.config.get("max_ddl_per_ticket", 5)
    if count > max_ddl:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(count=count, max=max_ddl),
        ))


# ── PG-DDL-012 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-012", name="PGDDLAllowColumnType",
    category=RuleCategory.DDL, severity=Severity.WARNING, db_type="postgres",
    description="ALTER TABLE ALTER TYPE 必须带 USING 子句",
    message_template="字段类型转换 ALTER TYPE 缺少 USING 子句，可能导致数据丢失",
))
def check_pg_ddl_012(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """允许字段类型转换（USING子句）。"""
    if not re.search(r"\bALTER\s+TABLE\b", sql_upper, re.IGNORECASE):
        return
    if re.search(r"\bALTER\s+(?:COLUMN\s+)?\w+\s+TYPE\b", sql_upper, re.IGNORECASE):
        if "USING" not in sql_upper:
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
            ))


# ── PG-DDL-013 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-013", name="PGAllowCreateView",
    category=RuleCategory.DDL, severity=Severity.INFO, db_type="postgres",
    description="检查 MATERIALIZED VIEW 的刷新策略",
    config={"allow_materialized_view": True, "refresh_strategy": "CONCURRENTLY"},
    message_template="创建 MATERIALIZED VIEW 时建议使用 REFRESH MATERIALIZED VIEW CONCURRENTLY",
))
def check_pg_ddl_013(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """允许创建视图（给出 MATERIALIZED VIEW 建议）。"""
    if not re.search(r"\bCREATE\s+(?:MATERIALIZED\s+)?VIEW\b", sql_upper, re.IGNORECASE):
        return
    if "MATERIALIZED" in sql_upper:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
        ))


# ── PG-DDL-014 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-014", name="PGDDLEnableForeignKey",
    category=RuleCategory.DDL, severity=Severity.WARNING, db_type="postgres",
    description="外键约束会导致共享行锁，影响并发性能",
    config={"allow_foreign_key": True, "warn_about_lock": True},
    message_template="外键约束(REFERENCES)在 PG 中可能导致共享行锁，影响并发写入性能，建议应用层控制",
))
def check_pg_ddl_014(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """允许外键（但警告PG外键有共享锁问题）。"""
    if not re.search(r"\bCREATE\s+TABLE\b|\bALTER\s+TABLE\b", sql_upper, re.IGNORECASE):
        return
    if re.search(r"\bREFERENCES\s+\w+", sql_upper, re.IGNORECASE) or "FOREIGN KEY" in sql_upper:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
        ))


# ── PG-DDL-015 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-015", name="PGAllowCreatePartition",
    category=RuleCategory.DDL, severity=Severity.INFO, db_type="postgres",
    description="检查分区表创建规范（Declarative Partitioning）",
    config={"supported_partition_types": ["RANGE", "LIST", "HASH"]},
    message_template="分区表建议使用 PARTITION BY {supported_types}，并确保有分区键索引",
))
def check_pg_ddl_015(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """允许创建分区表（Declarative Partitioning）。"""
    if not re.search(r"\bCREATE\s+TABLE\b", sql_upper, re.IGNORECASE):
        return
    if "PARTITION BY" in sql_upper:
        supported = rule.config.get("supported_partition_types", ["RANGE", "LIST", "HASH"])
        if not any(pt in sql_upper for pt in supported):
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                message=rule.message_template.format(supported_types="/".join(supported)),
            ))


# ── PG-DDL-016 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-016", name="PGDDLTablePrefix",
    category=RuleCategory.DDL, severity=Severity.INFO, db_type="postgres",
    description="表名前缀规范检查",
    config={"required_prefixes": []},
    message_template="表名 '{table}' 不符合前缀规范，建议前缀: {prefixes}",
))
def check_pg_ddl_016(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """表名前缀规范。"""
    required = rule.config.get("required_prefixes", [])
    if not required:
        return
    table_name = _extract_pg_table_name(sql_upper, sql_clean)
    if not table_name or not any(table_name.startswith(p) for p in required):
        if table_name:
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                message=rule.message_template.format(table=table_name, prefixes=", ".join(required)),
            ))


# ── PG-DDL-017 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-017", name="PGSupportEncoding",
    category=RuleCategory.DDL, severity=Severity.WARNING, db_type="postgres",
    description="数据库/表应使用 UTF8 编码",
    config={"allowed_encodings": ["UTF8"]},
    message_template="编码设置 '{encoding}' 不在推荐列表中，建议使用 UTF8",
))
def check_pg_ddl_017(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """允许的编码（UTF8推荐）。"""
    enc_match = re.search(r"\bENCODING\s*[=:]\s*['\"](\w+)['\"]|\bENCODING\s+(\w+)", sql_upper, re.IGNORECASE)
    if enc_match:
        encoding = enc_match.group(1) or enc_match.group(2)
        allowed = rule.config.get("allowed_encodings", ["UTF8"])
        if encoding and encoding.upper() not in [a.upper() for a in allowed]:
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                message=rule.message_template.format(encoding=encoding),
            ))


# ── PG-DDL-018 ──
@register_rule(AuditRuleV2(
    code="PG-DDL-018", name="PGMaxTableNameLen",
    category=RuleCategory.DDL, severity=Severity.WARNING, db_type="postgres",
    description="PostgreSQL 表名最大长度限制为 63 字节（NAMEDATALEN - 1）",
    config={"max_name_bytes": 63},
    message_template="表名 '{table}' 长度为 {length} 字节，超过 PG 限制 {max} 字节，超长部分将被截断",
))
def check_pg_ddl_018(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """表名最大长度（PG限制63字节）。"""
    table_name = _extract_pg_table_name(sql_upper, sql_clean)
    if not table_name:
        return
    name_bytes = len(table_name.encode("utf-8"))
    max_bytes = rule.config.get("max_name_bytes", 63)
    if name_bytes > max_bytes:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(table=table_name, length=name_bytes, max=max_bytes),
        ))


# ═══════════════════════════════════════════════
# DML 规则 (PG-DML-001 ~ PG-DML-008, 共8条)
# ═══════════════════════════════════════════════


# ── PG-DML-001 ──
@register_rule(AuditRuleV2(
    code="PG-DML-001", name="PGDMLMaxInsertRows",
    category=RuleCategory.DML, severity=Severity.WARNING, db_type="postgres",
    description="INSERT 语句最大行数限制",
    config={"max_insert_rows": 1000},
    message_template="单次 INSERT 包含 {count} 行，超过最大限制 {max}，请分批插入",
))
def check_pg_dml_001(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """Insert最大行数上限。"""
    if not re.search(r"\bINSERT\b", sql_upper, re.IGNORECASE):
        return
    max_rows = rule.config.get("max_insert_rows", 1000)
    values_match = re.search(r"\bVALUES\s+(.+)", sql_upper, re.IGNORECASE | re.DOTALL)
    if values_match:
        values_part = values_match.group(1)
        row_count = 0
        depth = 0
        for char in values_part:
            if char == "(":
                if depth == 0:
                    row_count += 1
                depth += 1
            elif char == ")":
                depth -= 1
        if row_count > max_rows:
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                message=rule.message_template.format(count=row_count, max=max_rows),
            ))


# ── PG-DML-002 ──
@register_rule(AuditRuleV2(
    code="PG-DML-002", name="PGDMLAllowLimitSTMT",
    category=RuleCategory.DML, severity=Severity.INFO, db_type="postgres",
    description="PostgreSQL 原生支持 UPDATE/DELETE 带 LIMIT，但需谨慎使用",
    config={"allow_limit": True, "warn_when_used": True},
    message_template="UPDATE/DELETE 使用了 LIMIT，可能导致部分数据未处理，建议配合 ORDER BY 使用",
))
def check_pg_dml_002(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """允许UPDATE/DELETE带LIMIT（PG原生支持，但给出警告）。"""
    for keyword in ("UPDATE", "DELETE"):
        if re.match(rf"\s*{keyword}\b", sql_upper, re.IGNORECASE):
            if re.search(r"\bLIMIT\s+\d+", sql_upper, re.IGNORECASE):
                if rule.config.get("warn_when_used", True):
                    report.findings.append(AuditFindingV2(
                        rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
                    ))
                    break


# ── PG-DML-003 ──
@register_rule(AuditRuleV2(
    code="PG-DML-003", name="PGDMLAllowInsertNull",
    category=RuleCategory.DML, severity=Severity.INFO, db_type="postgres",
    description="允许插入 NULL，但建议显式使用 DEFAULT",
    config={"allow_null": True, "prefer_default": True},
    message_template="INSERT 包含 NULL 值，建议改用 DEFAULT 关键字以保持语义清晰",
))
def check_pg_dml_003(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """允许插入Null（但建议使用DEFAULT）。"""
    if not re.search(r"\bINSERT\b", sql_upper, re.IGNORECASE):
        return
    if re.search(r"\bVALUES\b.*\bNULL\b", sql_upper, re.IGNORECASE):
        if rule.config.get("prefer_default", True):
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
            ))


# ── PG-DML-004 ──
@register_rule(AuditRuleV2(
    code="PG-DML-004", name="PGDMLWhere",
    category=RuleCategory.DML, severity=Severity.ERROR, db_type="postgres",
    description="UPDATE/DELETE 必须包含 WHERE 条件",
    message_template="{operation} 语句缺少 WHERE 条件，将操作全表数据，极其危险",
))
def check_pg_dml_004(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """强制DML必须拥有WHERE。"""
    for keyword in ("UPDATE", "DELETE"):
        if re.match(rf"\s*{keyword}\b", sql_upper, re.IGNORECASE):
            if "WHERE" not in sql_upper:
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                    message=rule.message_template.format(operation=keyword),
                ))


# ── PG-DML-005 ──
@register_rule(AuditRuleV2(
    code="PG-DML-005", name="PGDMLReturning",
    category=RuleCategory.DML, severity=Severity.INFO, db_type="postgres",
    description="推荐 DML 使用 RETURNING 子句获取操作结果",
    config={"recommend_returning": True},
    message_template="DML 语句建议使用 RETURNING 子句获取操作后的值（如 RETURNING id）",
))
def check_pg_dml_005(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """推荐DML使用RETURNING子句。"""
    for keyword in ("INSERT", "UPDATE", "DELETE"):
        if re.match(rf"\s*{keyword}\b", sql_upper, re.IGNORECASE):
            if "RETURNING" not in sql_upper:
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
                ))
                break


# ── PG-DML-006 ──
@register_rule(AuditRuleV2(
    code="PG-DML-006", name="PGDMLCTE",
    category=RuleCategory.DML, severity=Severity.INFO, db_type="postgres",
    description="CTE (Common Table Expression) 使用规范",
    config={"max_cte_depth": 3, "warn_recursive": True},
    message_template="CTE 使用规范: {cte_msg}",
))
def check_pg_dml_006(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """CTE使用规范。"""
    if not re.search(r"\bWITH\s+\w+\s+AS\s*\(", sql_upper, re.IGNORECASE):
        return
    if rule.config.get("warn_recursive", True) and re.search(r"\bWITH\s+RECURSIVE\b", sql_upper, re.IGNORECASE):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=Severity.WARNING,
            message="递归 CTE (WITH RECURSIVE) 可能导致无限循环，请确保有终止条件",
        ))
    cte_count = len(re.findall(r"\bWITH\b", sql_upper, re.IGNORECASE))
    max_depth = rule.config.get("max_cte_depth", 3)
    if cte_count > max_depth:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(cte_msg=f"检测到 {cte_count} 层嵌套 CTE，建议不超过 {max_depth} 层"),
        ))


# ── PG-DML-007 ──
@register_rule(AuditRuleV2(
    code="PG-DML-007", name="PGMaxAffectRows",
    category=RuleCategory.DML, severity=Severity.WARNING, db_type="postgres",
    description="DML 最大影响行数（通过 LIMIT 或 WHERE 范围控制）",
    config={"max_affected_rows": 10000},
    message_template="UPDATE/DELETE 缺少行数限制，可能影响超过 {max} 行数据，建议添加 LIMIT 或更精确的 WHERE",
))
def check_pg_dml_007(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """DML最大影响行数。"""
    max_rows = rule.config.get("max_affected_rows", 10000)
    for keyword in ("UPDATE", "DELETE"):
        if re.match(rf"\s*{keyword}\b", sql_upper, re.IGNORECASE):
            has_limit = re.search(r"\bLIMIT\s+\d+", sql_upper, re.IGNORECASE) is not None
            has_pk = re.search(r"\bid\s*[=<>]|\b\w+_id\s*[=<>]", sql_upper, re.IGNORECASE) is not None
            if not has_limit and not has_pk and "WHERE" in sql_upper:
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                    message=rule.message_template.format(max=max_rows),
                ))
            break


# ── PG-DML-008 ──
@register_rule(AuditRuleV2(
    code="PG-DML-008", name="PGDDLImplicitTypeConversion",
    category=RuleCategory.DML, severity=Severity.WARNING, db_type="postgres",
    description="禁止 WHERE 条件中的隐式类型转换，会导致索引失效",
    message_template="WHERE 条件中字段 '{column}' 与值之间存在隐式类型转换风险，请使用显式转换 ::TYPE",
))
def check_pg_dml_008(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """不允许隐式类型转换。"""
    if "WHERE" not in sql_upper:
        return
    where_match = re.search(r"\bWHERE\b(.+?)(?:GROUP|HAVING|ORDER|LIMIT|RETURNING|;|$)", sql_upper, re.IGNORECASE | re.DOTALL)
    if not where_match:
        return
    where_clause = where_match.group(1)
    patterns = [
        r"(\w+)\s*[=<>]\s*'\d+'",
        r"(\w+)\s*=\s*NULL",
    ]
    for pat in patterns:
        matches = re.findall(pat, where_clause, re.IGNORECASE)
        if matches:
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                message=rule.message_template.format(column=matches[0]),
            ))
            break


# ═══════════════════════════════════════════════
# DQL 规则 (PG-DQL-001 ~ PG-DQL-006, 共6条)
# ═══════════════════════════════════════════════


# ── PG-DQL-001 ──
@register_rule(AuditRuleV2(
    code="PG-DQL-001", name="PGDQLSelectStar",
    category=RuleCategory.DQL, severity=Severity.WARNING, db_type="postgres",
    description="SELECT * 返回不必要的列，影响网络传输和内存",
    message_template="禁止使用 SELECT *，请明确指定需要的字段（COUNT(*) 除外）",
))
def check_pg_dql_001(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """SELECT * 检测。"""
    if not re.search(r"\bSELECT\b", sql_upper, re.IGNORECASE):
        return
    if re.search(r"\bSELECT\s+\*\s+FROM\b", sql_upper, re.IGNORECASE):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
        ))


# ── PG-DQL-002 ──
@register_rule(AuditRuleV2(
    code="PG-DQL-002", name="PGDQLLimitOffset",
    category=RuleCategory.DQL, severity=Severity.INFO, db_type="postgres",
    description="大表分页建议使用游标分页替代 LIMIT/OFFSET",
    config={"max_offset_rows": 100000},
    message_template="使用 LIMIT/OFFSET 在大表上性能较差（OFFSET 越大越慢），建议改用游标分页 (WHERE id > ? LIMIT N) 或 keyset pagination",
))
def check_pg_dql_002(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """推荐LIMIT+OFFSET分页优化。"""
    if not re.search(r"\bSELECT\b", sql_upper, re.IGNORECASE):
        return
    if re.search(r"\bOFFSET\s+\d+\b", sql_upper, re.IGNORECASE):
        offset_match = re.search(r"\bOFFSET\s+(\d+)", sql_upper, re.IGNORECASE)
        if offset_match:
            offset_val = int(offset_match.group(1))
            max_offset = rule.config.get("max_offset_rows", 100000)
            if offset_val > max_offset:
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code, rule_name=rule.name, level=Severity.WARNING,
                    message=f"OFFSET {offset_val} 过大，超过建议值 {max_offset}，强烈建议改用游标分页",
                ))
            else:
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
                ))


# ── PG-DQL-003 ──
@register_rule(AuditRuleV2(
    code="PG-DQL-003", name="PGDQLWindowFunction",
    category=RuleCategory.DQL, severity=Severity.INFO, db_type="postgres",
    description="Window 函数使用检查，确保有合适的索引支持",
    config={"warn_no_partition": True},
    message_template="Window 函数 ({functions}) 使用 detected，建议确保 PARTITION BY / ORDER BY 列有索引支持",
))
def check_pg_dql_003(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """Window函数使用检查。"""
    window_funcs = [
        "ROW_NUMBER", "RANK", "DENSE_RANK", "NTILE",
        "LAG", "LEAD", "FIRST_VALUE", "LAST_VALUE",
        "NTH_VALUE", "PERCENT_RANK", "CUME_DIST",
    ]
    found_funcs = [f for f in window_funcs if re.search(rf"\b{f}\s*\(", sql_upper, re.IGNORECASE)]
    if found_funcs:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(functions=", ".join(found_funcs)),
        ))


# ── PG-DQL-004 ──
@register_rule(AuditRuleV2(
    code="PG-DQL-004", name="PGDQLRegex",
    category=RuleCategory.DQL, severity=Severity.WARNING, db_type="postgres",
    description="~ 操作符（正则表达式匹配）性能警告，无法使用索引",
    config={"prefer_prefix_regex": True},
    message_template="正则表达式匹配 (~ / ~*) 无法使用 B-Tree 索引，性能较差。建议: 1) 使用前缀匹配 ^ 配合 text_pattern_ops 索引; 2) 改用 pg_trgm 的 gin/gist 索引",
))
def check_pg_dql_004(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """~操作符性能警告。"""
    if re.search(r"\w+\s*~\*?\s*['\"]", sql_upper, re.IGNORECASE):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
        ))


# ── PG-DQL-005 ──
@register_rule(AuditRuleV2(
    code="PG-DQL-005", name="PGDQLConcurrentIndex",
    category=RuleCategory.DQL, severity=Severity.INFO, db_type="postgres",
    description="生产环境创建索引应使用 CREATE INDEX CONCURRENTLY 避免锁表",
    config={"require_concurrently": True},
    message_template="生产环境建议 CREATE INDEX CONCURRENTLY 代替 CREATE INDEX，避免阻塞写入",
))
def check_pg_dql_005(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """推荐CREATE INDEX CONCURRENTLY。"""
    if not re.search(r"\bCREATE\s+INDEX\b", sql_upper, re.IGNORECASE):
        return
    if "CONCURRENTLY" not in sql_upper:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
        ))


# ── PG-DQL-006 ──
@register_rule(AuditRuleV2(
    code="PG-DQL-006", name="PGDQLAnalyzeVacuum",
    category=RuleCategory.DQL, severity=Severity.INFO, db_type="postgres",
    description="大批量数据变更后应执行 ANALYZE 更新统计信息",
    config={"recommend_analyze": True},
    message_template="大批量数据操作后建议执行 ANALYZE table_name 更新统计信息，帮助优化器生成正确执行计划",
))
def check_pg_dql_006(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """检查ANALYZE更新统计信息。"""
    is_bulk = False
    if re.search(r"\bINSERT\b", sql_upper, re.IGNORECASE):
        values_match = re.search(r"\bVALUES\s+(.+)", sql_upper, re.IGNORECASE | re.DOTALL)
        if values_match:
            values_part = values_match.group(1)
            row_count = 0
            depth = 0
            for char in values_part:
                if char == "(":
                    if depth == 0:
                        row_count += 1
                    depth += 1
                elif char == ")":
                    depth -= 1
            if row_count > 100:
                is_bulk = True
    if re.search(r"\bUPDATE\s+\w+\s+SET\b", sql_upper, re.IGNORECASE) and "WHERE" in sql_upper:
        if not re.search(r"\bid\s*[=<>]", sql_upper, re.IGNORECASE) and "LIMIT" not in sql_upper:
            is_bulk = True
    if is_bulk:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity, message=rule.message_template,
        ))


# ═══════════════════════════════════════════════
# 安全规则 (PG-SEC-001 ~ PG-SEC-006, 共6条)
# ═══════════════════════════════════════════════


# ── PG-SEC-001 ──
@register_rule(AuditRuleV2(
    code="PG-SEC-001", name="PGSecuritySuperuser",
    category=RuleCategory.SECURITY, severity=Severity.ERROR, db_type="postgres",
    description="禁止使用超级用户(superuser)执行日常操作",
    config={"forbidden_roles": ["postgres", "rds_superuser"]},
    message_template="检测到超级用户操作 ({role})，禁止使用超级用户执行日常 SQL，应使用最小权限原则",
))
def check_pg_sec_001(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止超级用户操作。"""
    forbidden = rule.config.get("forbidden_roles", ["postgres", "rds_superuser"])
    superuser_patterns = [
        r"\bALTER\s+(?:ROLE|USER)\s+(\w+)\s+.*\bSUPERUSER\b",
        r"\bCREATE\s+(?:ROLE|USER)\s+(\w+).*\bSUPERUSER\b",
        r"\bSET\s+ROLE\s+(\w+)",
    ]
    for pat in superuser_patterns:
        match = re.search(pat, sql_upper, re.IGNORECASE)
        if match:
            role = match.group(1)
            if role in forbidden or "SUPERUSER" in sql_upper.upper():
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                    message=rule.message_template.format(role=role),
                ))
                break


# ── PG-SEC-002 ──
@register_rule(AuditRuleV2(
    code="PG-SEC-002", name="PGSecurityFileAccess",
    category=RuleCategory.SECURITY, severity=Severity.ERROR, db_type="postgres",
    description="禁止调用 pg_read_file, pg_read_binary_file 等文件访问函数",
    config={"forbidden_functions": [
        "pg_read_file", "pg_read_binary_file", "pg_ls_dir",
        "pg_stat_file", "pg_read_file_all",
    ]},
    message_template="禁止调用危险函数 '{func}'，该函数可访问服务器文件系统",
))
def check_pg_sec_002(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """禁止pg_read_file/pg_load。"""
    forbidden = rule.config.get("forbidden_functions", [])
    for func in forbidden:
        if re.search(rf"\b{func}\s*\(", sql_upper, re.IGNORECASE):
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                message=rule.message_template.format(func=func),
            ))
            break


# ── PG-SEC-003 ──
@register_rule(AuditRuleV2(
    code="PG-SEC-003", name="PGSecurityCopyFile",
    category=RuleCategory.SECURITY, severity=Severity.ERROR, db_type="postgres",
    description="COPY FROM/TO 可读写服务器文件系统，存在安全风险",
    config={"allow_copy": False, "allowed_paths": []},
    message_template="COPY {direction} 可访问文件系统，存在安全风险。如需使用请确保文件路径受控",
))
def check_pg_sec_003(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """COPY FROM/TO文件系统风险。"""
    if rule.config.get("allow_copy", False):
        return
    copy_from = re.search(r"\bCOPY\b.*\bFROM\b", sql_upper, re.IGNORECASE)
    copy_to = re.search(r"\bCOPY\b.*\bTO\b", sql_upper, re.IGNORECASE)
    if copy_from:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(direction="FROM"),
        ))
    if copy_to:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(direction="TO"),
        ))


# ── PG-SEC-004 ──
@register_rule(AuditRuleV2(
    code="PG-SEC-004", name="PGSecurityExtension",
    category=RuleCategory.SECURITY, severity=Severity.WARNING, db_type="postgres",
    description="危险扩展检查（dblink, plpython 等可执行外部代码）",
    config={
        "dangerous_extensions": [
            "dblink", "plpython3u", "plpython2u", "plperlu", "pltclu",
            "file_fdw", "postgres_fdw", "pg_file_settings",
            "adminpack", "auto_explain",
        ],
        "allowed_extensions": ["pgcrypto", "uuid-ossp", "pg_trgm"],
    },
    message_template="扩展 '{ext}' 存在安全风险，请评估必要性: 可访问外部资源或执行外部代码",
))
def check_pg_sec_004(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """危险扩展检查。"""
    if not re.search(r"\bCREATE\s+EXTENSION\b", sql_upper, re.IGNORECASE):
        return
    dangerous = rule.config.get("dangerous_extensions", [])
    allowed = rule.config.get("allowed_extensions", [])
    for ext in dangerous:
        if ext in allowed:
            continue
        if re.search(rf"\bCREATE\s+EXTENSION\s+(?:IF\s+NOT\s+EXISTS\s+)?{re.escape(ext)}\b", sql_upper, re.IGNORECASE):
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                message=rule.message_template.format(ext=ext),
            ))


# ── PG-SEC-005 ──
@register_rule(AuditRuleV2(
    code="PG-SEC-005", name="PGSecurityPassword",
    category=RuleCategory.SECURITY, severity=Severity.ERROR, db_type="postgres",
    description="密码复杂度检查",
    config={
        "min_length": 12,
        "require_special": True,
        "require_upper": True,
        "require_digit": True,
    },
    message_template="密码不符合复杂度要求: 最少 {min_len} 位，需包含大小写字母、数字和特殊字符",
))
def check_pg_sec_005(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """密码复杂度检查。"""
    pwd_match = re.search(
        r"\b(?:CREATE|ALTER)\s+(?:ROLE|USER)\s+\w+.*\bPASSWORD\s+['\"]([^'\"]+)['\"]",
        sql_clean, re.IGNORECASE,
    )
    if not pwd_match:
        return
    password = pwd_match.group(1)
    min_len = rule.config.get("min_length", 12)
    issues = []
    if len(password) < min_len:
        issues.append(f"长度不足 {min_len} 位")
    if rule.config.get("require_upper", True) and not re.search(r"[A-Z]", password):
        issues.append("缺少大写字母")
    if not re.search(r"[a-z]", password):
        issues.append("缺少小写字母")
    if rule.config.get("require_digit", True) and not re.search(r"\d", password):
        issues.append("缺少数字")
    if rule.config.get("require_special", True) and not re.search(r"[!@#$%^&*()_+\-=\[\]{};':\"|,.<>/?]", password):
        issues.append("缺少特殊字符")
    if issues:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=f"密码不符合复杂度要求: {', '.join(issues)}",
        ))


# ── PG-SEC-006 ──
@register_rule(AuditRuleV2(
    code="PG-SEC-006", name="PGSecurityListenNotify",
    category=RuleCategory.SECURITY, severity=Severity.INFO, db_type="postgres",
    description="LISTEN/NOTIFY 使用限制（防止 DoS 攻击）",
    config={"max_channels_per_session": 10, "warn_payload_size": 8000},
    message_template="LISTEN/NOTIFY 使用 detected: {msg}",
))
def check_pg_sec_006(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """LISTEN/NOTIFY使用限制。"""
    has_listen = re.search(r"\bLISTEN\s+\w+", sql_upper, re.IGNORECASE)
    has_notify = re.search(r"\bNOTIFY\s+\w+", sql_upper, re.IGNORECASE)
    if has_listen:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(
                msg="LISTEN 通道在会话断开时自动取消订阅，建议配置 UNLISTEN 清理"
            ),
        ))
    if has_notify:
        payload_match = re.search(r"\bNOTIFY\s+\w+\s*,\s*['\"]([^'\"]+)['\"]", sql_clean, re.IGNORECASE)
        if payload_match:
            payload = payload_match.group(1)
            max_size = rule.config.get("warn_payload_size", 8000)
            if len(payload.encode("utf-8")) > max_size:
                report.findings.append(AuditFindingV2(
                    rule_code=rule.code, rule_name=rule.name, level=Severity.WARNING,
                    message=f"NOTIFY payload 大小超过 {max_size} 字节（PG限制 8000 字节），建议减小消息体",
                ))


# ═══════════════════════════════════════════════
# 性能规则 (PG-PERF-001 ~ PG-PERF-007, 共7条)
# ═══════════════════════════════════════════════


# ── PG-PERF-001 ──
@register_rule(AuditRuleV2(
    code="PG-PERF-001", name="PGPerfVacuum",
    category=RuleCategory.PERFORMANCE, severity=Severity.INFO, db_type="postgres",
    description="大表 VACUUM 检查，建议避免 VACUUM FULL 阻塞操作",
    config={"warn_vacuum_full": True, "recommend_autovacuum": True, "table_size_threshold_mb": 1000},
    message_template="VACUUM 操作建议: {msg}",
))
def check_pg_perf_001(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """大表VACUUM检查。"""
    if not re.search(r"\bVACUUM\b", sql_upper, re.IGNORECASE):
        return
    if rule.config.get("warn_vacuum_full", True) and re.search(r"\bVACUUM\s+FULL\b", sql_upper, re.IGNORECASE):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=Severity.WARNING,
            message=rule.message_template.format(
                msg="VACUUM FULL 会获取排他锁并重建整个表，大表操作会长时间阻塞读写。建议优先使用标准 VACUUM 或 pg_repack"
            ),
        ))
    else:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(
                msg="建议启用 autovacuum 并合理配置阈值 (autovacuum_vacuum_scale_factor)，避免手动 VACUUM 遗漏"
            ),
        ))


# ── PG-PERF-002 ──
@register_rule(AuditRuleV2(
    code="PG-PERF-002", name="PGPerfSeqScan",
    category=RuleCategory.PERFORMANCE, severity=Severity.WARNING, db_type="postgres",
    description="检测可能导致全表扫描的查询模式",
    config={"min_table_size_for_warn": 10000},
    message_template="查询可能触发全表扫描: {reason}",
))
def check_pg_perf_002(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """全表扫描检测（静态检查）。"""
    if not re.search(r"\bSELECT\b", sql_upper, re.IGNORECASE):
        return
    reasons = []
    if "WHERE" not in sql_upper and "LIMIT" not in sql_upper:
        reasons.append("SELECT 缺少 WHERE 和 LIMIT 条件")
    where_match = re.search(r"\bWHERE\b(.+?)(?:GROUP|ORDER|LIMIT|OFFSET|;|$)", sql_upper, re.IGNORECASE | re.DOTALL)
    if where_match:
        where_clause = where_match.group(1)
        if re.search(r"\b\w+\s*\([^)]+\)\s*[=<>]", where_clause):
            reasons.append("WHERE 子句中列上使用了函数，可能导致索引失效")
        if re.search(r"\bLIKE\s+'%[^']+'", where_clause):
            reasons.append("LIKE 前导通配符 % 无法使用索引")
        if re.search(r"\b\w+\s+IS\s+NULL\b", where_clause):
            reasons.append("IS NULL 条件可能无法使用标准 B-Tree 索引")
    if re.search(r"\bILIKE\b", sql_upper, re.IGNORECASE):
        reasons.append("ILIKE 无法使用 B-Tree 索引，建议使用 citext 类型或 pg_trgm 扩展")
    if reasons:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(reason="; ".join(reasons)),
        ))


# ── PG-PERF-003 ──
@register_rule(AuditRuleV2(
    code="PG-PERF-003", name="PGPerfIndexBloat",
    category=RuleCategory.PERFORMANCE, severity=Severity.INFO, db_type="postgres",
    description="索引膨胀检测建议，REINDEX CONCURRENTLY 使用",
    config={"bloat_threshold_pct": 50, "recommend_reindex_concurrently": True},
    message_template="索引维护建议: {msg}",
))
def check_pg_perf_003(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """索引膨胀检测。"""
    if re.search(r"\bREINDEX\s+INDEX\b", sql_upper, re.IGNORECASE) and "CONCURRENTLY" not in sql_upper:
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=Severity.WARNING,
            message=rule.message_template.format(
                msg="REINDEX 会锁表，生产环境建议使用 REINDEX TABLE CONCURRENTLY 或 REINDEX INDEX CONCURRENTLY (PG 12+)"
            ),
        ))
    if re.search(r"\bCREATE\s+(?:UNIQUE\s+)?INDEX\b", sql_upper, re.IGNORECASE):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(
                msg="定期监控索引膨胀率 (pgstattuple 扩展)，膨胀超过 50% 时考虑 REINDEX CONCURRENTLY"
            ),
        ))


# ── PG-PERF-004 ──
@register_rule(AuditRuleV2(
    code="PG-PERF-004", name="PGPerfLockTimeout",
    category=RuleCategory.PERFORMANCE, severity=Severity.WARNING, db_type="postgres",
    description="锁超时配置检查，长时间持有锁影响并发",
    config={"recommended_lock_timeout_ms": 5000, "recommended_statement_timeout_ms": 30000},
    message_template="锁/超时配置建议: {msg}",
))
def check_pg_perf_004(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """锁超时配置。"""
    timeout_patterns = [
        (r"\bSET\s+lock_timeout\s*=\s*['\"]?(\d+\s*\w*)['\"]?", "lock_timeout"),
        (r"\bSET\s+statement_timeout\s*=\s*['\"]?(\d+\s*\w*)['\"]?", "statement_timeout"),
    ]
    for pat, name in timeout_patterns:
        match = re.search(pat, sql_upper, re.IGNORECASE)
        if match:
            value = match.group(1)
            report.findings.append(AuditFindingV2(
                rule_code=rule.code, rule_name=rule.name, level=rule.severity,
                message=rule.message_template.format(msg=f"{name} 设置为 {value}，建议合理配置避免长时间锁等待"),
            ))
    if re.search(r"\bUPDATE\s+\w+\s+SET\b", sql_upper, re.IGNORECASE) and "WHERE" in sql_upper and "LIMIT" not in sql_upper:
        recommended = rule.config.get("recommended_statement_timeout_ms", 30000)
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(
                msg=f"大批量操作建议在事务中设置 SET LOCAL statement_timeout = '{recommended}ms' 避免长时间持有锁"
            ),
        ))


# ── PG-PERF-005 ──
@register_rule(AuditRuleV2(
    code="PG-PERF-005", name="PGPerfConnectionPool",
    category=RuleCategory.PERFORMANCE, severity=Severity.INFO, db_type="postgres",
    description="连接池使用建议",
    config={"max_connections": 200, "recommended_pool_size": 20},
    message_template="连接管理建议: {msg}",
))
def check_pg_perf_005(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """连接池使用建议。"""
    if re.search(r"\bSET\s+(?:SESSION|LOCAL)\b", sql_upper, re.IGNORECASE):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(
                msg="频繁使用 SET SESSION/LOCAL 可能影响连接池复用，建议通过连接池统一配置参数"
            ),
        ))
    if re.search(r"\bBEGIN\b", sql_upper, re.IGNORECASE) and not re.search(r"\bCOMMIT\b|\bROLLBACK\b", sql_upper, re.IGNORECASE):
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=Severity.WARNING,
            message=rule.message_template.format(
                msg="检测到 BEGIN 但无 COMMIT/ROLLBACK，长时间打开的事务可能导致连接池耗尽和锁堆积"
            ),
        ))


# ── PG-PERF-006 ──
@register_rule(AuditRuleV2(
    code="PG-PERF-006", name="PGPerfWorkMem",
    category=RuleCategory.PERFORMANCE, severity=Severity.INFO, db_type="postgres",
    description="work_mem 配置建议，影响排序和哈希操作性能",
    config={"min_work_mem_mb": 4, "recommended_work_mem_mb": 16},
    message_template="内存配置建议: {msg}",
))
def check_pg_perf_006(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """work_mem配置建议。"""
    has_sort = False
    has_hash = False
    if "ORDER BY" in sql_upper:
        order_cols = re.findall(r"\bORDER\s+BY\s+(.+?)(?:LIMIT|OFFSET|;|$)", sql_upper, re.IGNORECASE)
        if order_cols and len(order_cols[0].split(",")) > 2:
            has_sort = True
    if "GROUP BY" in sql_upper:
        has_hash = True
    if re.search(r"\bDISTINCT\b", sql_upper, re.IGNORECASE):
        has_sort = True
    if "WITH" in sql_upper and ("JOIN" in sql_upper or "UNION" in sql_upper):
        has_sort = True
    if has_sort or has_hash:
        rec_mem = rule.config.get("recommended_work_mem_mb", 16)
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(
                msg=f"复杂排序/哈希操作建议确保 work_mem >= {rec_mem}MB，或设置 SET LOCAL work_mem = '{rec_mem}MB'"
            ),
        ))


# ── PG-PERF-007 ──
@register_rule(AuditRuleV2(
    code="PG-PERF-007", name="PGPerfParallelQuery",
    category=RuleCategory.PERFORMANCE, severity=Severity.INFO, db_type="postgres",
    description="并行查询使用规范（max_parallel_workers_per_gather）",
    config={"max_parallel_workers": 4, "min_table_size_for_parallel": "8MB"},
    message_template="并行查询建议: {msg}",
))
def check_pg_perf_007(sql_clean: str, sql_upper: str, report: AuditReportV2, rule: AuditRuleV2) -> None:
    """并行查询使用规范。"""
    if not re.search(r"\bSELECT\b", sql_upper, re.IGNORECASE):
        return
    is_complex = False
    join_count = len(re.findall(r"\bJOIN\b", sql_upper, re.IGNORECASE))
    if join_count >= 2:
        is_complex = True
    if re.search(r"\bUNION\s+ALL?\b", sql_upper, re.IGNORECASE):
        is_complex = True
    if re.search(r"\bCOUNT\s*\(|\bSUM\s*\(|\bAVG\s*\(", sql_upper, re.IGNORECASE) and "GROUP BY" in sql_upper:
        is_complex = True
    if is_complex:
        max_workers = rule.config.get("max_parallel_workers", 4)
        report.findings.append(AuditFindingV2(
            rule_code=rule.code, rule_name=rule.name, level=rule.severity,
            message=rule.message_template.format(
                msg=f"复杂查询可受益于并行查询，建议确认 max_parallel_workers_per_gather <= {max_workers}，并确保 min_parallel_table_scan_size 设置合理"
            ),
        ))


# ═══════════════════════════════════════════════
# 注册表导出
# ═══════════════════════════════════════════════

PG_RULE_REGISTRY: Dict[str, Tuple[AuditRuleV2, Callable]] = {}
"""PostgreSQL 特有规则注册表 (code -> (rule, check_func))。"""


def _init_pg_registry() -> None:
    """初始化 PostgreSQL 规则注册表。"""
    from app.services.audit_engine_v2 import _RULE_REGISTRY
    global PG_RULE_REGISTRY
    PG_RULE_REGISTRY = {
        code: (rule, func)
        for code, (rule, func) in _RULE_REGISTRY.items()
        if code.startswith("PG-")
    }


_init_pg_registry()


def get_pg_rules() -> List[AuditRuleV2]:
    """获取所有 PostgreSQL 审核规则定义。"""
    return sorted(
        [rule for rule, _ in PG_RULE_REGISTRY.values()],
        key=lambda r: r.code,
    )


def get_pg_rules_by_category(category: RuleCategory) -> List[AuditRuleV2]:
    """按分类获取 PostgreSQL 审核规则。"""
    return [rule for rule in get_pg_rules() if rule.category == category]


def get_pg_rule_count() -> Dict[str, int]:
    """获取 PostgreSQL 规则统计信息。"""
    rules = get_pg_rules()
    return {
        "total": len(rules),
        "ddl": sum(1 for r in rules if r.category == RuleCategory.DDL),
        "dml": sum(1 for r in rules if r.category == RuleCategory.DML),
        "dql": sum(1 for r in rules if r.category == RuleCategory.DQL),
        "security": sum(1 for r in rules if r.category == RuleCategory.SECURITY),
        "performance": sum(1 for r in rules if r.category == RuleCategory.PERFORMANCE),
    }


# ──────────────────────────────────────────────
# 模块加载自检
# ──────────────────────────────────────────────
if __name__ == "__main__":
    stats = get_pg_rule_count()
    print("=" * 60)
    print("PostgreSQL 审核规则引擎 - 规则统计")
    print("=" * 60)
    print(f"DDL 规则:        {stats['ddl']} 条 (PG-DDL-001 ~ PG-DDL-018)")
    print(f"DML 规则:        {stats['dml']} 条 (PG-DML-001 ~ PG-DML-008)")
    print(f"DQL 规则:        {stats['dql']} 条 (PG-DQL-001 ~ PG-DQL-006)")
    print(f"安全规则:        {stats['security']} 条 (PG-SEC-001 ~ PG-SEC-006)")
    print(f"性能规则:        {stats['performance']} 条 (PG-PERF-001 ~ PG-PERF-007)")
    print("-" * 60)
    print(f"总计:            {stats['total']} 条规则")
    print("=" * 60)
    print("\n规则明细:")
    for rule in get_pg_rules():
        print(f"  {rule.code} | {rule.severity.value:7s} | {rule.name}")