from pathlib import Path
import shutil

base = Path(__file__).parent
for root in (Path('C:/wechat_rpa'), Path('C:/kefu/wechat_rpa')):
    path = root / 'wxwork_message_browser.py'
    raw = path.read_bytes()
    source = raw.decode('utf-8').replace('\r\n', '\n')
    before = '''        rows = _read_message_rows(message_path, index, [index["conversations"][key]["lastRowId"] for key in missing])
        for key in missing:
            item = rows.get(index["conversations"][key]["lastRowId"], {})
            index["previews"][key] = {
                "text": _content_label(item.get("content"), item.get("content_type")).replace("\\n", " ")[:2000],
                "sender": str(item.get("sender_id") or ""),
            }
'''
    after = '''        for offset in range(0, len(missing), 500):
            batch = missing[offset:offset + 500]
            rows = _read_message_rows(message_path, index, [index["conversations"][key]["lastRowId"] for key in batch])
            for key in batch:
                item = rows.get(index["conversations"][key]["lastRowId"], {})
                index["previews"][key] = {
                    "text": _content_label(item.get("content"), item.get("content_type")).replace("\\n", " ")[:2000],
                    "sender": str(item.get("sender_id") or ""),
                }
'''
    assert source.count(before) == 1
    source = source.replace(before, after)
    path.write_bytes(source.replace('\n', '\r\n' if b'\r\n' in raw else '\n').encode('utf-8'))

test = base / 'test_wxwork_browser_performance.py'
source = test.read_text(encoding='utf-8')
assert 'test_concurrent_account_switches_cannot_leak' not in source
source = source.replace("\n\nif __name__ == '__main__':", '\n' + (base / 'browser_extra_tests.txt').read_text(encoding='utf-8') + "\n\nif __name__ == '__main__':")
test.write_text(source, encoding='utf-8')
for root in (Path('C:/wechat_rpa'), Path('C:/kefu/wechat_rpa')):
    shutil.copy2(test, root / test.name)