import ast
from pathlib import Path
import shutil


def edit_method(text, name, transform):
    node = next(n for n in ast.walk(ast.parse(text)) if isinstance(n, ast.FunctionDef) and n.name == name)
    lines = text.splitlines(keepends=True)
    old = ''.join(lines[node.lineno - 1:node.end_lineno])
    lines[node.lineno - 1:node.end_lineno] = [transform(old)]
    return ''.join(lines)


for root in [Path(r'C:\wechat_rpa'), Path(r'C:\kefu\wechat_rpa')]:
    backup = root / 'backups' / 'review-badges-20260917'
    for name in ('test_visual_review_visibility.py', 'test_reply_safety.py', 'test_protocol_review_visibility.py'):
        path = root / name
        if not path.exists():
            continue
        if not (backup / name).exists():
            shutil.copy2(path, backup / name)
        text = path.read_text(encoding='utf-8-sig')
        if name == 'test_visual_review_visibility.py':
            text = text.replace("'reply_text': '合成待审草稿', 'judge_risk': 'high',",
                                "'reply_text': '合成待审草稿', 'judge_risk': 'high', 'judge_mode': 'score_only',")
            text = text.replace("        self.bot.send_mode = 'auto'\n", "        self.bot.send_mode = 'auto'\n        self.enterContext(mock.patch('review_policy.current_judge_mode', return_value='score_only'))\n", 1)
        elif name == 'test_protocol_review_visibility.py':
            text = edit_method(text, 'test_high_risk_reason_available_in_automatic_mode', lambda s: s.replace(
                "update(judge_risk='high')", "update(judge_risk='high',judge_mode='score_only')").replace(
                '        state = self.enqueue()', "        with mock.patch('review_policy.current_judge_mode',return_value='score_only'):\n            state = self.enqueue()"))
        else:
            for method in ('test_high_judge_risk_triggers_review_even_with_no_rules', 'test_judge_risk_check_is_case_insensitive_and_takes_priority'):
                text = edit_method(text, method, lambda s: s.replace('with mock.patch.dict(', 'with mock.patch("review_policy.current_judge_mode", return_value="score_only"), mock.patch.dict(', 1).replace(
                    '{"judge_risk": "high"}', '{"judge_risk": "high", "judge_mode": "score_only"}').replace(
                    '{"judge_risk": "HIGH"}', '{"judge_risk": "HIGH", "judge_mode": "score_only"}').replace(
                    '裁判判高风险这一条不在配置里、关不掉——一条规则都没配也要拦。', '非影子裁判判高风险时，一条关键词规则都没配也要拦。'))
            text = edit_method(text, 'test_auto_mode_holds_a_reply_the_judge_scored_high_risk', lambda s: s.replace(
                '"judge_risk": "high",', '"judge_risk": "high", "judge_mode": "score_only",').replace(
                '        with (\n', '        with (\n            mock.patch("review_policy.current_judge_mode", return_value="score_only"),\n', 1))
        compile(text, str(path), 'exec')
        temporary = path.with_name(path.name + '.review-policy.tmp')
        temporary.write_text(text, encoding='utf-8')
        temporary.replace(path)
print('Updated old high-risk tests to explicitly cover enforcing judge mode.')