#!/usr/bin/env python3
# -*- coding: UTF-8 -*-
'''
@Project :domainScanDemo
@File    :create_sensitive_words_table.py
@IDE     :PyCharm
@Author  :梦伴
@Date    :2026/4/10 03:30
@explain : 创建敏感词表
'''

import psycopg2
from app.config import config


def create_sensitive_words_table():
    """
    创建敏感词表
    """
    try:
        # 连接数据库
        conn = psycopg2.connect(
            host=config.DB_HOST,
            port=config.DB_PORT,
            database=config.DB_DATABASE,
            user=config.DB_USER,
            password=config.DB_PASSWORD
        )
        print(f"成功连接到数据库: {config.DB_HOST}:{config.DB_PORT}/{config.DB_DATABASE}")
        
        # 创建游标
        cur = conn.cursor()
        
        # 创建敏感词表
        create_table_sql = """
        CREATE TABLE IF NOT EXISTS sensitive_words (
            id SERIAL PRIMARY KEY,
            word VARCHAR(255) UNIQUE NOT NULL,
            category VARCHAR(50) DEFAULT 'default',
            priority INTEGER DEFAULT 1,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        """
        cur.execute(create_table_sql)
        print("创建敏感词表成功")
        
        # 添加表注释
        add_table_comment_sql = """
        COMMENT ON TABLE sensitive_words IS '敏感词表，存储需要过滤的敏感词'
        """
        cur.execute(add_table_comment_sql)
        print("添加表注释成功")
        
        # 添加字段注释
        add_column_comments_sql = """
        COMMENT ON COLUMN sensitive_words.id IS '主键ID';
        COMMENT ON COLUMN sensitive_words.word IS '敏感词';
        COMMENT ON COLUMN sensitive_words.category IS '敏感词分类';
        COMMENT ON COLUMN sensitive_words.priority IS '优先级，数字越大优先级越高';
        COMMENT ON COLUMN sensitive_words.created_at IS '创建时间';
        COMMENT ON COLUMN sensitive_words.updated_at IS '更新时间';
        """
        cur.execute(add_column_comments_sql)
        print("添加字段注释成功")
        
        # 创建索引
        create_index_sql = """
        CREATE INDEX IF NOT EXISTS idx_sensitive_words_word ON sensitive_words(word);
        CREATE INDEX IF NOT EXISTS idx_sensitive_words_category ON sensitive_words(category);
        """
        cur.execute(create_index_sql)
        print("创建索引成功")
        
        # 创建更新触发器
        create_function_sql = """
        CREATE OR REPLACE FUNCTION update_updated_at_column()
        RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at = CURRENT_TIMESTAMP;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;
        """
        cur.execute(create_function_sql)
        print("创建更新函数成功")
        
        # 检查触发器是否存在
        check_trigger_sql = """
        SELECT COUNT(*) FROM pg_trigger WHERE tgname = 'update_sensitive_words_updated_at'
        """
        cur.execute(check_trigger_sql)
        trigger_exists = cur.fetchone()[0] > 0
        
        if not trigger_exists:
            create_trigger_sql = """
            CREATE TRIGGER update_sensitive_words_updated_at
            BEFORE UPDATE ON sensitive_words
            FOR EACH ROW
            EXECUTE FUNCTION update_updated_at_column();
            """
            cur.execute(create_trigger_sql)
            print("创建触发器成功")
        else:
            print("触发器已存在，跳过创建")
        
        # 提交事务
        conn.commit()
        print("事务提交成功")
        
        # 关闭游标和连接
        cur.close()
        conn.close()
        print("连接关闭成功")
        
        return True
    except Exception as e:
        print(f"创建敏感词表失败: {e}")
        return False


if __name__ == "__main__":
    print("开始创建敏感词表...")
    success = create_sensitive_words_table()
    if success:
        print("敏感词表创建成功！")
    else:
        print("敏感词表创建失败！")