from pathlib import Path
import shutil

base = Path(__file__).parent
helpers = (base / 'browser_cache_helpers.txt').read_text(encoding='utf-8')
helpers = helpers.replace('"names": 16, "files": 16, "messages": 32', '"names": 64, "files": 64, "counts": 64, "messages": 32')
helpers = helpers.replace('conversation_id.startswith("Y:")', 'conversation_id.upper().startswith("Y:")')
helpers += (base / 'browser_counts.txt').read_text(encoding='utf-8')
queries = (base / 'browser_queries.txt').read_text(encoding='utf-8')

for root in (Path('C:/wechat_rpa'), Path('C:/kefu/wechat_rpa')):
    path = root / 'wxwork_message_browser.py'
    original = path.read_bytes()
    source = original.decode('utf-8').replace('\r\n', '\n')
    def one(before, after):
        global source
        assert source.count(before) == 1, (before[:100], source.count(before))
        source = source.replace(before, after)
    one('import sqlite3\n', 'import sqlite3\nimport math\nimport threading\nfrom array import array\nfrom collections import OrderedDict\n')
    one('def _format_time(value: Any) -> str:\n', helpers + 'def _format_time(value: Any) -> str:\n')
    one('def _load_names(databases: dict[str, Path], account: str) -> tuple[dict[str, str], dict[str, str], str]:', 'def _load_names_uncached(databases: dict[str, Path], account: str) -> tuple[dict[str, str], dict[str, str], str]:')
    one('def _database_files(databases: dict[str, Path]) -> tuple[list[dict[str, Any]], int]:', 'def _database_files_uncached(databases: dict[str, Path]) -> tuple[list[dict[str, Any]], int]:')
    one('                    row_count = _count(connection, table)\n', '                    row_count = _message_counts(path)[0] if table == "message_table" else _count(connection, table)\n')
    start = source.index('def _conversation_rows(')
    end = source.index('def load_browser_snapshot(', start)
    source = source[:start] + queries + source[end:]
    one('    message_limit: int = 500,\n', '    message_limit: int = 500,\n    conversation_offset: int = 0,\n    message_offset: int = 0,\n')
    one('        discovered = _discover_databases(roots)\n', '        clear_browser_cache()\n        discovered = _discover_databases(roots)\n')
    start = source.index('        message_count = 0\n')
    end = source.index('        updated = max(\n', start)
    source = source[:start] + '''        message_count = 0
        conversation_count = 0
        message_path = databases.get("message.db")
        if message_path:
            try:
                message_count, conversation_count = _message_counts(message_path)
            except (sqlite3.Error, TypeError, ValueError):
                pass
''' + source[end:]
    one('            max(1, min(int(conversation_limit), 1000)),\n', '            max(1, min(int(conversation_limit), 1000)),\n            offset=max(0, int(conversation_offset)),\n')
    one('            message_limit,\n', '            message_limit,\n            offset=max(0, int(message_offset)),\n')
    backup = root / 'backups/loading-performance-20260917'
    backup.mkdir(parents=True, exist_ok=True)
    assert not (backup / path.name).exists(), 'Already patched, inspect first'
    shutil.copy2(path, backup / path.name)
    path.write_bytes(source.replace('\n', '\r\n' if b'\r\n' in original else '\n').encode('utf-8'))
    print(root, 'browser cache/index applied')

# sqlite connection context managers commit/rollback but do not close.
benchmark = base / 'benchmark_browser.py'
source = benchmark.read_text(encoding='utf-8')
source = source.replace('    original = browser.connect_sqlite\n', '    db.close()\n    original = browser.connect_sqlite\n')
benchmark.write_text(source, encoding='utf-8')