from pathlib import Path
import shutil

methods = Path(__file__).with_name('visual_handoff_methods.txt').read_text(encoding='utf-8')


def one(source, before, after):
    assert source.count(before) == 1, (before[:100], source.count(before))
    return source.replace(before, after)


for root in (Path('C:/wechat_rpa'), Path('C:/kefu/wechat_rpa')):
    path = root / 'wechat_bot.py'
    raw = path.read_bytes()
    source = raw.decode('utf-8').replace('\r\n', '\n')
    source = one(source, '    def cancel_pending_replies(self, keys) -> dict:\n', methods + '\n    def cancel_pending_replies(self, keys) -> dict:\n')
    source = one(source, '''        return bool(key and key in getattr(self, "_cancelled_reply_sessions", set()))
''', '''        return bool(key and (
            key in getattr(self, "_cancelled_reply_sessions", set())
            or self._manual_handoff_key(fp)
        ))
''')
    source = one(source, '''        path = getattr(self, "_pending_reply_path", "")
        if not path or not os.path.exists(path):
''', '''        self._manual_handoffs = {}
        path = getattr(self, "_pending_reply_path", "")
        if not path or not os.path.exists(path):
''')
    source = one(source, '''        restored = {}
        now = time.time()
        for key, value in raw.items():
''', '''        handoffs = raw.get("manual_handoffs")
        if isinstance(handoffs, dict):
            self._manual_handoffs = {str(key): dict(entry) for key, entry in handoffs.items()
                                     if isinstance(entry, dict) and isinstance(entry.get("state"), dict)}
        raw = raw["pending"] if isinstance(raw.get("pending"), dict) else raw
        restored = {}
        now = time.time()
        for key, value in raw.items():
''')
    source = one(source, '''            restored[fp.hex()] = state
''', '''            if not self._manual_handoff_key(fp):
                restored[fp.hex()] = state
''')
    source = one(source, '''        tmp = path + ".tmp"
        try:
            with open(tmp, "w", encoding="utf-8") as handle:
                json.dump(serializable, handle, ensure_ascii=False, indent=2)
''', '''        handoffs = getattr(self, "_manual_handoffs", {})
        if handoffs:
            serializable = {"pending": serializable, "manual_handoffs": handoffs}
        tmp = path + ".tmp"
        try:
            with open(tmp, "w", encoding="utf-8") as handle:
                json.dump(serializable, handle, ensure_ascii=False, indent=2)
''')
    marker = '                if fp.hex() in getattr(self, "_cancelled_reply_sessions", set()):\n'
    assert source.count(marker) == 2
    source = source.replace(marker, '                if self._reply_task_cancelled(fp):\n')
    source = one(source, '''            with self._pending_lock:
                if self._reply_task_cancelled(fp):
''', '''            with self._pending_lock:
                if not self._resume_handoff_for_new_message(fp, database_event=database_event):
                    return False, "本会话已转人工，尚未证明是后续新消息"
                if self._reply_task_cancelled(fp):
''')
    source = one(source, '''        if key in getattr(self, "_cancelled_reply_sessions", set()):
            return False
        newly_queued = key not in pending
''', '''        if self._reply_task_cancelled(fp):
            self._notice_manual_handoff(fp)
            return False
        newly_queued = key not in pending
''')
    source = one(source, '''        database_direction = self._database_message_direction(fp, chat_text)
        if database_direction is not None:
''', '''        if not self._resume_handoff_for_new_message(fp, str(chat_text or "")):
            return False
        database_direction = self._database_message_direction(fp, chat_text)
        if database_direction is not None:
''')
    source = one(source, '''            recovered_text = "\n".join(recovered_delta)
'''.replace('"\n"', '"\\n"'), '''            self._resume_handoff_for_new_message(fp, str(startup_chat_text or ""))
            recovered_text = "\\n".join(recovered_delta)
''')
    source = one(source, '''            if pending_saved is False:
                print("    [待回复恢复] 未读任务无法落盘，已禁止点击以避免丢失消息。")
''', '''            if pending_saved is False:
                if self._manual_handoff_key(target_fp):
                    self._notice_manual_handoff(target_fp)
                    continue
                print("    [待回复恢复] 未读任务无法落盘，已禁止点击以避免丢失消息。")
''')
    backup = root / 'backups/review-visibility-20260916/wechat_bot.before-handoff.py'
    assert not backup.exists(), 'Already applied; inspect first'
    shutil.copy2(path, backup)
    path.write_bytes(source.replace('\n', '\r\n' if b'\r\n' in raw else '\n').encode('utf-8'))
    print(root, 'visual handoff applied')