"""Idempotent source patches and CPU checks. No vLLM/torch imports or XPU needed."""
import argparse
import ast
import copy
import hashlib
import inspect
import json
from importlib import metadata
from pathlib import Path
from types import SimpleNamespace

from routing_compat import patch_source, regression_tests, INSERT

parser = argparse.ArgumentParser()
parser.add_argument('--site-packages', type=Path,
                    default=Path('/opt/vllm_kunlun/lib/python3.10/site-packages'))
parser.add_argument('--check', action='store_true', help='Verify without changing installed files')
args = parser.parse_args()
root = args.site_packages
route = root / 'vllm_kunlun/quantization/compressed_tensors/compressed_tensors_moe.py'
wrapper = root / 'vllm_kunlun/ops/_custom_ops.py'
lower = root / 'kunlun_ops/_xflashinfer_moe.py'
versions = {p: metadata.version(p) for p in ('vllm', 'vllm-kunlun', 'kunlun-ops')}
if not versions['vllm'].startswith('0.25.1') or not versions['vllm-kunlun'].startswith('0.25.1'):
    raise RuntimeError(f'Unsupported environment: {versions}')

before_route = route.read_text()
after_route = patch_source(before_route)
# Reproduce the legacy bug and exercise patched method dispatch.
regression_tests(after_route.replace(INSERT, ''), after_route)
lower_fn = next(n for n in ast.parse(lower.read_text()).body
                if isinstance(n, ast.FunctionDef) and n.name == 'moe_sigmoid_group_topk_norm')
signature = inspect.Signature([inspect.Parameter(a.arg, inspect.Parameter.POSITIONAL_OR_KEYWORD)
                               for a in lower_fn.args.args])
assert 'block_statistic' in signature.parameters

before_wrapper = wrapper.read_text()
lines = before_wrapper.splitlines(keepends=True)
names = {'moe_sigmoid_group_topk_norm', 'moe_sigmoid_group_topk_norm_cuda'}
functions = [n for n in ast.parse(before_wrapper).body if isinstance(n, ast.FunctionDef) and n.name in names]
assert len(functions) == 2
for fn in functions:
    for call in ast.walk(fn):
        if isinstance(call, ast.Call) and ast.unparse(call.func) == 'kunlun_ops.moe_sigmoid_group_topk_norm':
            for kw in call.keywords:
                if kw.arg == 'block_static':
                    idx = kw.lineno - 1
                    assert 'block_static=block_static' in lines[idx]
                    lines[idx] = lines[idx].replace('block_static=block_static', 'block_statistic=block_static')
after_wrapper = ''.join(lines)
for fn in [n for n in ast.parse(after_wrapper).body if isinstance(n, ast.FunctionDef) and n.name in names]:
    node = copy.deepcopy(fn)
    node.decorator_list = []; node.returns = None
    for a in node.args.args:
        a.annotation = None
    calls = []
    def spy(*a, **kw):
        calls.append(signature.bind(*a, **kw).arguments)
    scope = {'kunlun_ops': SimpleNamespace(moe_sigmoid_group_topk_norm=spy)}
    exec(compile(ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[])), '<wrapper-check>', 'exec'), scope)
    values = {k: object() for k in ('x', 'topk_index', 'norm_score', 'block_static', 'bias')}
    scope[fn.name](**values, scale=1.0, n_group=8, topk_group=4)
    assert len(calls) == 1 and calls[0]['block_statistic'] is values['block_static']
    assert calls[0]['bias'] is values['bias']
    print('WRAPPER_PASS', fn.name)

for path, before, after in [(route, before_route, after_route), (wrapper, before_wrapper, after_wrapper)]:
    compile(after, str(path), 'exec')
    if args.check:
        if before != after:
            raise RuntimeError(f'Missing patch: {path}')
    elif before != after:
        backup = path.with_name(path.name + '.bak-image-build')
        if backup.exists():
            raise RuntimeError(f'Backup already exists: {backup}')
        backup.write_text(before)
        temp = path.with_name(path.name + '.image-build.tmp')
        temp.write_text(after); temp.chmod(path.stat().st_mode & 0o777); temp.replace(path)
    print('SHA256', hashlib.sha256(after.encode()).hexdigest(), path)
print('VERSIONS', json.dumps(versions))
print('ALL_PATCH_CHECKS_PASS')