"""补充更新 special_types 和 number_features（基于标题重算）"""
import re
import yaml
from datetime import datetime

def load_config():
    with open('/root/coolbot-data/config.yaml', 'r') as f:
        return yaml.safe_load(f)

config = load_config()
DB_CONFIG = config['database']

def get_db_conn():
    import psycopg2
    return psycopg2.connect(
        host=DB_CONFIG['host'], port=DB_CONFIG.get('port', 5432),
        user=DB_CONFIG['user'], password=DB_CONFIG['password'],
        database=DB_CONFIG['database']
    )

FEATURE_MAP = {
    '倒置': ['倒置'], '如意': ['如意'], '朦胧': ['朦胧'], '金马': ['金马'], '金山': ['金山'],
    '天马': ['天马'], '钻石': ['钻石'], '永恒': ['永恒'], '无4': ['无4', '无四'],
    '无47': ['无47'], '无247': ['无247'], '无34': ['无34'], '无347': ['无347'],
    '带4': ['带4'], '豹子': ['豹子'], '狮子': ['狮子'], '老虎': ['老虎'],
    '大象': ['大象'], '生日': ['生日'], '满号': ['满号'], '首日': ['首日'],
}

SPECIAL_MAP = {
    '标十': ['标十'], '标百': ['标百'],
    '刀': ['刀', '刀货'], '单张': ['单张', '散张'],
    '捆': ['捆'], '千连': ['千连'], '百连': ['百连'],
    '救生圈': ['大救生圈', '救生圈'],
}

def calc_special(title):
    if not title: return None
    found = [k for k, words in SPECIAL_MAP.items() if any(w in title for w in words)]
    return '|'.join(found) if found else None

def calc_features(title):
    if not title: return None
    found = [k for k, words in FEATURE_MAP.items() if any(w in title for w in words)]
    return '|'.join(found) if found else None

conn = get_db_conn()
cur = conn.cursor()
cur.execute("SELECT id, title FROM yichens_posts WHERE special_types IS NULL OR number_features IS NULL")
rows = cur.fetchall()
print(f'需要更新: {len(rows)} 条')

updated = 0
for rid, title in rows:
    sp = calc_special(title)
    ft = calc_features(title)
    cur.execute("""
        UPDATE yichens_posts SET special_types = COALESCE(%s, special_types),
        number_features = COALESCE(%s, number_features)
        WHERE id = %s
    """, (sp, ft, rid))
    updated += 1
    if updated % 200 == 0:
        print(f'已更新 {updated}/{len(rows)}')

conn.commit()
cur.close()
conn.close()
print(f'完成! 共更新 {updated} 条')