"""Run inside srx_test: fix lower-level keyword without changing Torch schema."""
import ast
import copy
import datetime
import inspect
from pathlib import Path
from types import SimpleNamespace

import kunlun_ops

path = Path('/opt/vllm_kunlun/lib/python3.10/site-packages/vllm_kunlun/ops/_custom_ops.py')
original = path.read_text()
signature = inspect.signature(kunlun_ops.moe_sigmoid_group_topk_norm)
assert 'block_statistic' in signature.parameters
names = {'moe_sigmoid_group_topk_norm', 'moe_sigmoid_group_topk_norm_cuda'}
tree = ast.parse(original)
functions = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in names]
assert len(functions) == 2
lines = original.splitlines(keepends=True)
changed = 0
for fn in functions:
    for call in ast.walk(fn):
        if not isinstance(call, ast.Call) or ast.unparse(call.func) != 'kunlun_ops.moe_sigmoid_group_topk_norm':
            continue
        for keyword in call.keywords:
            if keyword.arg == 'block_static':
                index = keyword.lineno - 1
                assert 'block_static=block_static' in lines[index]
                lines[index] = lines[index].replace('block_static=block_static', 'block_statistic=block_static')
                changed += 1
assert changed in (0, 2), f'Unexpected matching calls: {changed}'
patched = ''.join(lines)
compile(patched, str(path), 'exec')
# Execute both actual wrapper bodies against the installed lower-level signature.
for fn in [n for n in ast.parse(patched).body if isinstance(n, ast.FunctionDef) and n.name in names]:
    node = copy.deepcopy(fn)
    node.decorator_list = []
    node.returns = None
    for arg in node.args.args:
        arg.annotation = None
    calls = []
    def spy(*args, **kwargs):
        bound = signature.bind(*args, **kwargs)
        calls.append(bound.arguments)
    scope = {'kunlun_ops': SimpleNamespace(moe_sigmoid_group_topk_norm=spy)}
    exec(compile(ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[])), '<wrapper-test>', 'exec'), scope)
    sentinels = {name: object() for name in ('x', 'topk_index', 'norm_score', 'block_static', 'bias')}
    scope[fn.name](**sentinels, scale=1.0, n_group=8, topk_group=4)
    assert calls[0]['block_statistic'] is sentinels['block_static']
    assert calls[0]['bias'] is sentinels['bias']
    print('WRAPPER_SIGNATURE_PASS', fn.name)
if changed:
    backup = path.with_name(path.name + '.bak-sigmoid-keyword-' + datetime.datetime.now().strftime('%Y%m%dT%H%M%S%f'))
    backup.write_text(original)
    temp = path.with_name(path.name + '.routing-fix.tmp')
    temp.write_text(patched)
    temp.chmod(path.stat().st_mode & 0o777)
    temp.replace(path)
    print('BACKUP', backup)
    print('PATCHED', path)
else:
    print('ALREADY_PATCHED')