"""Make YuE's HF token samplers honor Comfy's cancel flag. Run with the node directory."""
import ast
import pathlib
import shutil
import sys

CHECK = '        from comfy.model_management import throw_exception_if_processing_interrupted\n        throw_exception_if_processing_interrupted()\n'

def patch(source):
    tree = ast.parse(source)
    cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == 'BlockTokenRangeProcessor')
    method = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == '__call__')
    if 'throw_exception_if_processing_interrupted' in ast.get_source_segment(source, method):
        return source
    lines = source.splitlines(keepends=True)
    lines.insert(method.body[0].lineno - 1, CHECK)
    result = ''.join(lines)
    ast.parse(result)
    return result

if __name__ == '__main__':
    root = pathlib.Path(sys.argv[1]).resolve(strict=True)
    plans = []
    for name in ('infer.py', 'common.py'):
        path = root / 'inference' / name
        source = path.read_text(encoding='utf-8')
        plans.append((path, source, patch(source)))
    for path, source, result in plans:
        if source == result:
            print(f'Already patched: {path.name}')
            continue
        backup = path.with_suffix('.py.before-aigen-cancel')
        if backup.exists():
            raise RuntimeError(f'Backup already exists: {backup}')
        shutil.copy2(path, backup)
        path.write_text(result, encoding='utf-8')
        print(f'Patched {path.name}; original backed up.')