import os
from lawrisk.utils.env_loader import load_env
from lawrisk.services.licensing_repo import _lic_pg_conn

def add_test_rule():
    load_env()
    
    with _lic_pg_conn(autocommit=True) as conn:
        cur = conn.cursor()
        
        # 1. Ensure "测试主题" exists
        cur.execute("INSERT INTO themes (name) VALUES (%s) ON CONFLICT (name) DO UPDATE SET name = EXCLUDED.name RETURNING id", ("测试主题",))
        theme_id = cur.fetchone()[0]
        print(f"Theme '测试主题' ID: {theme_id}")
        
        # 2. Add rule for "测试许可" -> "测试主题"
        # Check if already exists
        cur.execute("SELECT id FROM permit_theme_rules WHERE permit_name = %s AND theme_id = %s", ("测试许可", theme_id))
        if cur.fetchone():
            print("Rule '测试许可' -> '测试主题' already exists")
        else:
            cur.execute("INSERT INTO permit_theme_rules (permit_name, theme_id) VALUES (%s, %s)", ("测试许可", theme_id))
            print("Rule '测试许可' -> '测试主题' added successfully")

if __name__ == "__main__":
    add_test_rule()