#!/usr/bin/env python3
"""Patch the paused/running container's installed Kunlun INT8 MoE, with backup.

Run on the Docker host. Tests exercise the actual method through routing dispatch
with tensor allocation/kernel calls mocked; they do not execute XPU kernels.
"""
import ast
import copy
import datetime
import hashlib
import io
import subprocess
import tarfile
from types import SimpleNamespace as NS

CONTAINER = 'srx_test'
DIRECTORY = '/opt/vllm_kunlun/lib/python3.10/site-packages/vllm_kunlun/quantization/compressed_tensors'
NAME = 'compressed_tensors_moe.py'
CLASS = 'KunlunCompressedTensorsW8A8Int8MoEMethod'
MARKER = '# Kunlun vLLM 0.25 routing compatibility fix'
INSERT = '''        # Kunlun vLLM 0.25 routing compatibility fix
        # RoutedExperts.forward_monolithic passes only layer/x/logits/input_ids.
        # Read routing from the layer, not legacy method defaults. In particular,
        # routed_scaling_factor may be 1 because the runner applies it externally.
        scoring_func = layer.scoring_func
        num_expert_group = layer.num_expert_group
        topk_group = layer.topk_group
        e_score_correction_bias = layer.e_score_correction_bias
        routed_scaling_factor = layer.routed_scaling_factor
        if scoring_func not in ("softmax", "sigmoid"):
            raise NotImplementedError(f"Unsupported Kunlun INT8 MoE routing: {scoring_func}")

'''


def method_node(source):
    module = ast.parse(source)
    cls = next(n for n in module.body if isinstance(n, ast.ClassDef) and n.name == CLASS)
    return next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == 'apply_monolithic')


def patch_source(source):
    if MARKER in source:
        return source
    node = method_node(source)
    first = node.body[0]
    if ast.unparse(first) != 'hidden_states = x':
        raise RuntimeError('Unexpected installed method; refusing to patch')
    lines = source.splitlines(keepends=True)
    lines.insert(first.lineno - 1, INSERT)
    patched = ''.join(lines)
    compile(patched, NAME, 'exec')
    return patched


class Captured(Exception):
    pass


def dispatch(source, scoring, groups, top_groups, scale, bias):
    node = copy.deepcopy(method_node(source))
    for arg in node.args.posonlyargs + node.args.args + node.args.kwonlyargs:
        arg.annotation = None
    node.returns = None
    calls = []

    def capture(kind):
        def callback(**kwargs):
            calls.append((kind, kwargs))
            raise Captured()
        return callback

    ops = NS(moe_softmax_topk_norm=capture('softmax'),
             moe_sigmoid_group_topk_norm=capture('sigmoid'))
    torch = NS(float32='float32', int32='int32',
               empty=lambda *a, **k: object(), zeros=lambda *a, **k: object(), ops=NS(_C=ops))
    scope = {'torch': torch}
    exec(compile(ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[])), '<actual-method>', 'exec'), scope)
    layer = NS(scoring_func=scoring, num_expert_group=groups, topk_group=top_groups,
               routed_scaling_factor=scale, e_score_correction_bias=bias,
               w13_weight=NS(shape=(256, 512, 128)), w2_weight=NS(shape=(256, 128, 256)))
    try:
        scope['apply_monolithic'](NS(moe=NS(experts_per_token=8)), layer,
                                  NS(shape=(1, 128), device='cpu'), NS(float=lambda: object()),
                                  input_ids=None)
    except Captured:
        pass
    assert len(calls) == 1
    return calls[0]


def regression_tests(original, patched):
    bias = object()
    if MARKER not in original:
        kind, _ = dispatch(original, 'sigmoid', 8, 4, 1.0, bias)
        assert kind == 'softmax', 'Expected to reproduce the old routing bug'
        print('REPRODUCED: old method used softmax for DeepSeek sigmoid routing')
    for scale in (1.0, 2.5):
        kind, kwargs = dispatch(patched, 'sigmoid', 8, 4, scale, bias)
        assert kind == 'sigmoid'
        assert kwargs['n_group'] == 8 and kwargs['topk_group'] == 4
        assert kwargs['bias'] is bias and kwargs['scale'] == scale
    assert dispatch(patched, 'softmax', None, None, 1.0, None)[0] == 'softmax'
    kind, kwargs = dispatch(patched, 'sigmoid', None, None, 1.0, bias)
    assert kind == 'sigmoid' and kwargs['bias'] is bias
    print('PASS: sigmoid groups/bias, external/internal scaling, softmax and ungrouped sigmoid')


def read_installed():
    data = subprocess.check_output(['docker', 'cp', f'{CONTAINER}:{DIRECTORY}/{NAME}', '-'])
    with tarfile.open(fileobj=io.BytesIO(data)) as archive:
        entry = next(m for m in archive.getmembers() if m.isfile())
        return entry, archive.extractfile(entry).read()


def main():
    entry, original_bytes = read_installed()
    original = original_bytes.decode()
    patched = patch_source(original)
    regression_tests(original, patched)
    if patched == original:
        print('ALREADY_PATCHED')
        return
    backup = NAME + '.bak-routing-' + datetime.datetime.now().strftime('%Y%m%dT%H%M%S%f')
    output = io.BytesIO()
    with tarfile.open(fileobj=output, mode='w') as archive:
        for name, content in [(backup, original_bytes), (NAME, patched.encode())]:
            info = copy.copy(entry)
            info.name, info.size = name, len(content)
            archive.addfile(info, io.BytesIO(content))
    subprocess.run(['docker', 'cp', '-', f'{CONTAINER}:{DIRECTORY}'], input=output.getvalue(), check=True)
    _, installed = read_installed()
    assert installed == patched.encode(), 'Installed content differs from patch'
    print('INSTALLED', DIRECTORY + '/' + NAME)
    print('BACKUP', DIRECTORY + '/' + backup)
    print('SHA256', hashlib.sha256(installed).hexdigest())
    print('Restart the vLLM service to load the fix; existing workers retain old code.')


if __name__ == '__main__':
    main()