"""
Patches transformers 4.55.3 to register the qwen3_5 model type.

Deploy steps on the remote machine:
  1. cp -r modified_scripts/qwen3_5 /usr/local/lib/python3.10/site-packages/transformers/models/qwen3_5
  2. python3 modified_scripts/patch_transformers_qwen3_5.py

Target: pip-installed transformers at /usr/local/lib/python3.10/site-packages/transformers/
(Not the corex pre-installed path at /usr/local/corex/lib64/python3/dist-packages/)
"""

import sys

TRANSFORMERS_ROOT = "/usr/local/lib/python3.10/site-packages/transformers"
AUTO_CONFIG = f"{TRANSFORMERS_ROOT}/models/auto/configuration_auto.py"
MODELS_INIT = f"{TRANSFORMERS_ROOT}/models/__init__.py"


def patch_file(path, replacements):
    with open(path, "r") as f:
        content = f.read()

    patched = False
    for old, new in replacements:
        if new in content:
            print(f"  [skip] already patched: {repr(new[:60])}")
            continue
        if old not in content:
            print(f"  [warn] anchor not found: {repr(old[:60])}")
            continue
        content = content.replace(old, new, 1)
        patched = True
        print(f"  [ok]   inserted after: {repr(old[:60])}")

    if patched:
        with open(path, "w") as f:
            f.write(content)


def main():
    print(f"=== Patching {AUTO_CONFIG} ===")
    patch_file(AUTO_CONFIG, [
        # CONFIG_MAPPING_NAMES: insert qwen3_5 right after qwen3
        (
            '("qwen3", "Qwen3Config"),',
            '("qwen3", "Qwen3Config"),\n        ("qwen3_5", "Qwen3_5Config"),',
        ),
        # Some versions don't have trailing comma — handle that too
        (
            '("qwen3", "Qwen3Config")\n',
            '("qwen3", "Qwen3Config"),\n        ("qwen3_5", "Qwen3_5Config"),\n',
        ),
        # MODEL_NAMES_MAPPING (model_type -> human readable name, used by docstring generator)
        (
            '("qwen3", "Qwen3"),',
            '("qwen3", "Qwen3"),\n        ("qwen3_5", "Qwen3_5"),',
        ),
        (
            '("qwen3", "Qwen3")\n',
            '("qwen3", "Qwen3"),\n        ("qwen3_5", "Qwen3_5"),\n',
        ),
    ])

    print(f"\n=== Patching {MODELS_INIT} ===")
    patch_file(MODELS_INIT, [
        (
            "from .qwen3 import *\n",
            "from .qwen3 import *\n    from .qwen3_5 import *\n",
        ),
    ])

    # Verification
    print("\n=== Verification ===")
    try:
        import importlib.util, types

        # Quick smoke-test: import the config class directly
        spec = importlib.util.spec_from_file_location(
            "configuration_qwen3_5",
            f"{TRANSFORMERS_ROOT}/models/qwen3_5/configuration_qwen3_5.py",
        )
        mod = importlib.util.module_from_spec(spec)
        # Provide minimal parent package stubs so relative imports resolve
        pkg = types.ModuleType("transformers")
        pkg.__path__ = [TRANSFORMERS_ROOT]
        sys.modules.setdefault("transformers", pkg)
        spec.loader.exec_module(mod)
        cfg = mod.Qwen3_5Config()
        print(f"  Qwen3_5Config() smoke-test OK  (model_type={cfg.model_type})")
    except Exception as e:
        print(f"  [warn] smoke-test failed (may be fine at runtime): {e}")

    print("\nDone.")


if __name__ == "__main__":
    main()