from pathlib import Path
import shutil

root = Path(r'C:\wechat_rpa')
target = root / 'protocol_engine.py'
backup = Path(__file__).parent / 'protocol-before'
backup.mkdir(parents=True, exist_ok=True)
if not (backup / target.name).exists():
    shutil.copy2(target, backup / target.name)
source = target.read_text(encoding='utf-8')
source = source.replace("reply_text='',log=None,progress=None):", "reply_text='',log=None,progress=None,review_assistant=None):", 1)
source = source.replace("        self._lock=threading.RLock();self._owned_db=db is None", "        self._lock=threading.RLock();self._owned_db=db is None\n        self._review_assistant=review_assistant;self._assistant_warning=''", 1)
anchor = '    def _save(self):'
helpers = Path(__file__).with_name('protocol_assistant_helpers.txt').read_text(encoding='utf-8')
assert anchor in source and 'def _assistant_call(' not in source
source = source.replace(anchor, helpers + '\n' + anchor, 1)
anchor = "        phase={'queued':'等待消息合并','generating':'正在生成回复','awaiting_review':'等待审核',"
source = source.replace(anchor, "        if stage=='awaiting_review':self._register_assistant_review(key,s)\n" + anchor, 1)
anchor = "        if self._skip_excluded_session(key,message):return"
source = source.replace(anchor, """        if self._assistant_recipient(message['conv_id']):
            with self._lock:
                self._pending.pop(key,None)
                if event:self._seen[key]=event
                self._save()
            return
        if old and message.get('is_self'):
            context=self.db.get_conversation_context_by_id(self.account,message['conv_id'],limit=500)
            if self._reconcile_assistant_review(key,old,context):return
            if self._assistant_pending(old):return
""" + anchor, 1)
source = source.replace("'detected_by':'protocol_db','detection_ts':now,'reply_text':'','awaiting_review':False,'content_type':message.get('content_type')}", "'detected_by':'protocol_db','detection_ts':now,'reply_text':'','awaiting_review':False,'content_type':message.get('content_type'),\n                'database_event':dict(message)}", 1)
anchor = '        self._completed_polls+=1'
source = source.replace(anchor, "        self._assistant_call('enqueue_notifications',default=None)\n" + anchor, 1)
anchor = "        if self._skip_excluded_session(key,state,expected_state=state):\n            if excluded is not None:excluded.append(_session_exclusion(state))"
source = source.replace(anchor, """        if self._assistant_recipient(state['conv_id']):
            if excluded is not None:excluded.append('配置的医疗助理会话不参与自动回复')
            return False
""" + anchor, 1)
anchor = "        last=(context or {}).get('last_message')\n        if not last:return False"
source = source.replace(anchor, """        if self._reconcile_assistant_review(key,state,context):return False
        last=(context or {}).get('last_message')
        if not last:return False
        if last.get('is_self') and self._assistant_pending(state):return False""", 1)
anchor = "        if _send_protection(state):return\n        if self._skip_excluded_session(key,state,expected_state=state):return"
source = source.replace(anchor, """        if _send_protection(state):return
        if self._assistant_recipient(state['conv_id']):
            with self._lock:
                if self._pending.get(key) is state:self._pending.pop(key,None);self._save()
            return
        if self._reconcile_assistant_review(key,state):return
""" + '        if self._skip_excluded_session(key,state,expected_state=state):return', 1)
start = source.index("        with self._lock:\n            if self._pending.get(key) is not state:return\n            state.update(send_state='sending',awaiting_review=False)")
end = source.index('    def _needs_review(self,state):', start)
send_body = source[start:end]
send_body = send_body.replace('        self.reply_count+=1', "        self._assistant_call('record_automated_send',state,receipt,default=None)\n        self.reply_count+=1", 1)
source = source[:start] + '''        import send_lock
        if not send_lock.try_acquire('protocol:'+str(threading.get_ident()),timeout=0):return
        try:
            if not self._current(key,state):return
            self._send_ready(key,state)
        finally:send_lock.release()

    def _send_ready(self,key,state):
''' + send_body + source[end:]
old = '''    def _needs_review(self,state):
        if self.send_mode=='review':return True,'人工审核后点击“通过并发送”'
        import ai_config
        from review_policy import selective_review_reason
        reason=selective_review_reason(state,state.get('staged_user_text',''),state.get('reply_text',''),
                                       rules=getattr(ai_config,'AI_REVIEW_RULES',[]) or [])
        return bool(reason),reason
'''
new = '''    def _needs_review(self,state):
        import ai_config
        from review_policy import selective_review_reason
        reason=selective_review_reason(state,state.get('staged_user_text',''),state.get('reply_text',''),
                                       rules=getattr(ai_config,'AI_REVIEW_RULES',[]) or [])
        if reason:return True,reason
        if self._assistant_pending(state):
            return True,state.get('review_reason') or '本会话仍有待处理审核，请由医疗助理回复客户'
        if self.send_mode=='review' and not self._assistant_call('resumed',state.get('conv_id','')):
            return True,'人工审核后点击“通过并发送”'
        return False,''
'''
assert old in source
source = source.replace(old,new,1)
source = source.replace('    def close(self):\n        self.sender.close()', "    def close(self):\n        if self._review_assistant is not None:\n            try:self._review_assistant.close()\n            except Exception:pass\n        self.sender.close()", 1)
target.write_text(source,encoding='utf-8',newline='')
print('protocol assistant integration patched')