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 = 'wanted = [name for name in ("sender_id", "send_time", "content_type", "content", "server_id", "client_id") if name in columns]'
    assert source.count(before) == 1
    source = source.replace(before, 'wanted = [name for name in ("conversation_id", "sender_id", "send_time", "content_type", "content", "server_id", "client_id") if name in columns]')
    before = '''                item = rows.get(index["conversations"][key]["lastRowId"], {})
                index["previews"][key] = {
'''
    after = '''                item = rows.get(index["conversations"][key]["lastRowId"], {})
                # A cache file can be atomically replaced between indexing and
                # reading. Reused row IDs must never expose another chat.
                if str(item.get("conversation_id") or "") != key:
                    index["previews"][key] = {"text": "", "sender": ""}
                    continue
                index["previews"][key] = {
'''
    assert source.count(before) == 1
    source = source.replace(before, after)
    before = '''        if item is None:
            continue
        sender_id = str(item.get("sender_id") or "")
'''
    after = '''        if item is None or str(item.get("conversation_id") or "") != conversation_id:
            continue
        sender_id = 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_database_replaced_between_index' not in source
source = source.replace("\n\nif __name__ == '__main__':", '\n' + (base / 'browser_race_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)