from pathlib import Path
import hashlib
import json

roots = [Path('C:/kefu/wechat_rpa'), Path('C:/wechat_rpa')]
for root in roots:
    path = root / 'test_ai_context_audit.py'
    text = path.read_text(encoding='utf-8')
    marker = '    def test_dify_single_user_query_stays_unchanged(self):'
    addition = '''    def test_dify_generic_completion_preserves_judge_system_and_history(self):
        messages = [{'role': 'system', 'content': '裁判必须输出JSON，参考知识KB-742。'}, *self.history_as_messages(), {'role': 'user', 'content': '那周六呢？'}]
        provider = self.provider('dify')
        with mock.patch('ai_chat._call_dify', return_value='{"score": 90}') as call:
            ai_chat._chat_completion(messages, provider=provider)
        self.assertIn('必须输出JSON', call.call_args.args[0])
        self.assertIn('KB-742', call.call_args.args[0])
        self.assertIn('蓝色方案', call.call_args.args[0])
        self.assertIs(call.call_args.kwargs['provider'], provider)

'''
    assert marker in text
    text = text.replace(marker, addition + marker, 1)
    text = text.replace('AI_MCP_ENABLED=False, AI_USE_VISION=False))', 'AI_MCP_ENABLED=False, AI_USE_VISION=False, AI_DEVELOPMENT_MODE=False))', 1)
    path.write_text(text, encoding='utf-8')

    path = root / 'ai_chat.py'
    raw = path.read_bytes()
    text = raw.decode('utf-8').replace('\r\n', '\n')
    old = '''        last_user = ""
        for m in reversed(messages):
            if m.get("role") == "user":
                c = m.get("content")
                last_user = c if isinstance(c, str) else str(c)
                break
        answer = _call_dify(last_user or "请回复", provider=provider)'''
    assert old in text
    text = text.replace(old, '        answer = _call_dify(model_protocol.dify_query(messages), provider=provider)', 1)
    newline = '\r\n' if b'\r\n' in raw else '\n'
    path.write_bytes(text.replace('\n', newline).encode('utf-8'))

manifest_path = Path('C:/kefu/deploy/recognition-audit-20260916/ai-context-manifest.json')
manifest = json.loads(manifest_path.read_text(encoding='utf-8'))
for entry in manifest:
    entry['after_sha256'] = hashlib.sha256(Path(entry['path']).read_bytes()).hexdigest()
manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding='utf-8')
print('Completed generic Dify context preservation and one additional regression.')