from pathlib import Path
import os
import shutil

wrapper = '''

_key_acquisition_lock = threading.Lock()
_key_acquisition_inflight = {}


def acquire_local_keys(timeout: float = 60.0) -> None:
    """Share concurrent initialization of the same local data directory only.

    There is no completed-result cache: a later explicit request scans again.
    Every caller still validates its own account database after acquisition.
    """
    root = application_data_dir()
    try:
        config = json.loads((root / "wxwork_gui_config.json").read_text(encoding="utf-8-sig"))
        source = str(config.get("db_dir") or "") if isinstance(config, dict) else ""
    except (OSError, ValueError):
        source = ""
    identity = (str(root.resolve()).casefold(), source.casefold(),
                os.environ.get("WECOM_ARCHIVE_SOURCE_DIR", "").casefold())
    with _key_acquisition_lock:
        future = _key_acquisition_inflight.get(identity)
        owner = future is None
        if owner:
            future = Future()
            _key_acquisition_inflight[identity] = future
    if not owner:
        try:
            return future.result(timeout=max(0.0, float(timeout)))
        except FutureTimeoutError as exc:
            raise ValueError("正在获取本机密钥，等待超时。请稍后重试，或导入已有密钥文件") from exc
    try:
        result = _acquire_local_keys_once(timeout)
    except BaseException as exc:
        future.set_exception(exc)
        raise
    else:
        future.set_result(result)
        return result
    finally:
        with _key_acquisition_lock:
            if _key_acquisition_inflight.get(identity) is future:
                _key_acquisition_inflight.pop(identity, None)
'''

for root in (Path(r'C:\kefu\wechat_rpa'), Path(r'C:\wechat_rpa')):
    path = root / 'wxwork_local_setup.py'
    backup = root / 'backups/loading-performance-20260917'
    backup.mkdir(parents=True, exist_ok=True)
    assert not (backup / path.name).exists()
    shutil.copy2(path, backup / path.name)
    text = path.read_text(encoding='utf-8')
    text = text.replace('import tempfile\n', 'import tempfile\nimport threading\nfrom concurrent.futures import Future, TimeoutError as FutureTimeoutError\n', 1)
    text = text.replace('def acquire_local_keys(timeout: float = 60.0) -> None:', 'def _acquire_local_keys_once(timeout: float = 60.0) -> None:', 1)
    text += wrapper
    compile(text, str(path), 'exec')
    tmp = path.with_suffix('.performance-tmp');tmp.write_text(text, encoding='utf-8');os.replace(tmp, path)
    print(str(root) + ': simultaneous key requests share one worker')