#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
test_wrapper_e2e.py —— 用平台真实命令形态端到端验证 vllm_wrapper.sh

背景：平台 bi-100 × vllm × text-generation 的 build-config 实际下发的是
    vllm serve /model --port 80 --served-model-name llm --max-model-len 4096 \
        --gpu-memory-utilization 0.9 --enforce-eager --trust-remote-code -tp 1
因为 Dockerfile 把镜像里的 vllm 换成了本仓 wrapper，这条命令会流经 wrapper。
本测试用一个"假 vllm_real"（把收到的参数原样打印）验证：
  1. wrapper 正确识别 serve 形态，不吞参数、不串位
  2. preflight 产出的修补参数被**追加**在原命令之后（不覆盖、不前置）
  3. 非 serve 子命令（如 --version）原样透传
  4. preflight 自己挂掉时不阻断启动（仍能 exec vllm_real）
  5. 模型目录里有真实缺陷时，参数确实被补上

不需要 GPU / vLLM / Docker。
"""
import io
import json
import os
import shutil
import subprocess
import sys
import tempfile

HERE = os.path.dirname(os.path.abspath(__file__))
WRAPPER = os.path.join(HERE, 'vllm_wrapper.sh')
PREFLIGHT = os.path.join(HERE, 'preflight.py')

PASS, FAIL = [], []


def check(name, cond, detail=''):
    (PASS if cond else FAIL).append(name)
    print(('  PASS ' if cond else '  FAIL ') + name + (('  -- ' + str(detail)) if (detail and not cond) else ''))


def make_fake_preflight(dirpath, extra_args, fail=False):
    """造一个假 preflight：按给定 extra_args 输出 JSON；fail=True 时非 0 退出。"""
    p = os.path.join(dirpath, 'preflight.py')
    body = [
        'import json, sys',
        'args = sys.argv[1:]',
        'out = ""',
        'if "--out" in args: out = args[args.index("--out") + 1]',
        'if ' + ('True' if fail else 'False') + ':',
        '    sys.stderr.write("boom' + chr(92) + 'n"); sys.exit(3)',
        'payload = {"extra_args": ' + repr(extra_args) + ', "patches": ["stub"], "overlay": ""}',
        'if out:',
        '    fh = open(out, "w")',
        '    fh.write(json.dumps(payload))',
        '    fh.close()',
        'print(json.dumps(payload))',
    ]
    with io.open(p, 'w', encoding='utf-8', newline=chr(10)) as f:
        f.write(chr(10).join(body) + chr(10))
    return p


def run_wrapper(model_dir, args, box, extra_args=None, fail_preflight=False):
    """在隔离目录跑 wrapper。

    box = 隔离目录；里面放：vllm_real（假）、preflight.py（假）、wrapper_local.sh。
    extra_args 由假 preflight 产出，用来验证 wrapper 的参数拼装。
    全部使用相对路径（cwd=box），规避 Windows 路径被 MSYS 吃掉。
    """
    real = os.path.join(box, 'vllm_real')
    with io.open(real, 'w', encoding='utf-8', newline=chr(10)) as f:
        f.write('#!/bin/bash' + chr(10) + 'echo "VLLM_REAL_ARGS: $@"' + chr(10))
    os.chmod(real, 0o755)

    make_fake_preflight(box, extra_args or [], fail=fail_preflight)

    src = io.open(WRAPPER, encoding='utf-8', newline=chr(10)).read()
    local = os.path.join(box, 'wrapper_local.sh')
    with io.open(local, 'w', encoding='utf-8', newline=chr(10)) as f:
        f.write(src
                .replace('/usr/local/corex/lib64/python3/dist-packages/bin/vllm_real', './vllm_real')
                .replace('/opt/preflight.py', './preflight.py')
                .replace('/tmp/mhxc_preflight.json', './preflight.json')
                .replace('/tmp/mhxc_preflight.out', './preflight.out')
                .replace('/tmp/mhxc_preflight.err', './preflight.err'))
    os.chmod(local, 0o755)

    env = dict(os.environ)
    env['PATH'] = box + os.pathsep + env.get('PATH', '')
    cmd = ['bash', 'wrapper_local.sh'] + args
    p = subprocess.run(cmd, capture_output=True, text=True, timeout=90, env=env, cwd=box)
    return p.returncode, p.stdout, p.stderr


def parse_real_args(stdout):
    for l in stdout.splitlines():
        if l.startswith('VLLM_REAL_ARGS: '):
            return l[len('VLLM_REAL_ARGS: '):]
    return None


def main():
    tmp = tempfile.mkdtemp(prefix='mhxc_e2e_')
    try:
        PLATFORM_CMD = ['serve', '/model', '--port', '80', '--served-model-name', 'llm',
                        '--max-model-len', '4096', '--gpu-memory-utilization', '0.9',
                        '--enforce-eager', '--trust-remote-code', '-tp', '1']
        EXTRA = ['--tokenizer', '/tmp/ov/tokenizer',
                 '--chat-template', '/tmp/ov/tpl.jinja',
                 '--hf-overrides', '{"architectures":["Qwen3_5MoeForCausalLM"]}']

        print('[1] 平台真实命令 + preflight 产出三类补丁')
        rc, out, err = run_wrapper('/model', PLATFORM_CMD, tmp, extra_args=EXTRA)
        final = parse_real_args(out) or ''
        check('wrapper 正常退出', rc == 0, err[-300:])
        check('真 vllm 被调起', final != '', out[:200])
        orig = PLATFORM_CMD[2:]
        check('原命令参数全部保留', all(a in final.split() for a in orig), final)
        check('serve 与 /model 在最前', final.split()[:2] == ['serve', '/model'], final[:100])
        check('三类补丁都带上', all(a in final.split() for a in ['--tokenizer', '--chat-template', '--hf-overrides']), final)
        # 正确性要求不是"补丁必须在末尾"，而是"平台原命令的参数序列不被拆乱"：
        # 平台参数应以原有顺序连续出现，且补丁只作为追加项存在。
        toks = final.split()
        plat = PLATFORM_CMD[2:]
        letidx = toks.index(plat[0]) if plat[0] in toks else -1
        seq_ok = letidx >= 0
        if seq_ok:
            for k in range(len(plat)):
                if toks[letidx + k] != plat[k]:
                    seq_ok = False
                    break
        check('平台原命令参数序列完整且顺序不变', seq_ok, (letidx, final))
        head_ok = toks[0] == 'serve' and toks[1] == '/model'
        check('serve 与模型目录仍在最前（位置参数不被挤掉）', head_ok, final[:80])
        print('     最终命令: ' + final[:260])

        print('[2] 干净模型：preflight 无产出 → 命令与平台原始命令完全一致')
        rc2, out2, err2 = run_wrapper('/model', PLATFORM_CMD, tmp, extra_args=[])
        final2 = parse_real_args(out2) or ''
        check('无补丁时命令逐字一致', final2.strip() == ' '.join(PLATFORM_CMD).strip(), final2)

        print('[3] 非 serve 子命令原样透传')
        rc3, out3, err3 = run_wrapper('/model', ['--version'], tmp, extra_args=EXTRA)
        check('--version 原样透传', (parse_real_args(out3) or '').strip() == '--version', out3[:200])

        print('[4] preflight 挂掉时不阻断启动')
        rc4, out4, err4 = run_wrapper('/model', PLATFORM_CMD, tmp, fail_preflight=True)
        final4 = parse_real_args(out4) or ''
        check('仍能起服务且用原命令', rc4 == 0 and final4.strip() == ' '.join(PLATFORM_CMD).strip(), (rc4, final4[:150]))

        print('[5] preflight 自己写 JSON 也能被正确读回（不依赖 --out 文件时的兜底）')
        rc5, out5, err5 = run_wrapper('/model', PLATFORM_CMD, tmp, extra_args=['--dtype', 'half'])
        check('单参数补丁也能拼上', '--dtype' in (parse_real_args(out5) or '').split(), out5[:200])

        print('[6] bash 语法')
        p = subprocess.run(['bash', '-n', 'vllm_wrapper.sh'], capture_output=True, text=True, cwd=HERE)
        check('wrapper bash 语法正确', p.returncode == 0, p.stderr[:200])

        print(chr(10) + '===== 结果: %d passed, %d failed =====' % (len(PASS), len(FAIL)))
        if FAIL:
            print('失败项: ' + ', '.join(FAIL))
            return 1
        return 0
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


if __name__ == '__main__':
    sys.exit(main())