"""Start Comfy with HF-provided custom nodes while leaving the image unchanged."""
import atexit,json,os,subprocess,sys,time,urllib.request
from pathlib import Path
PROCESS=None

def stop():
    global PROCESS
    if PROCESS and PROCESS.poll() is None:
        PROCESS.terminate()
        try:PROCESS.wait(timeout=3)
        except subprocess.TimeoutExpired:PROCESS.kill()
    PROCESS=None

def extra_paths(root):
    paths={'base_path':str(root),'is_default':True}
    if (root/'custom_nodes').is_dir():paths['custom_nodes']='custom_nodes'
    if (root/'models').is_dir():
        for folder in (root/'models').iterdir():
            if folder.is_dir() and folder.name.replace('_','').isalnum():paths[folder.name]='models/'+folder.name
    return paths

def initialize():
    global PROCESS
    value=os.environ.get('HF_APP_ROOT')
    if not value:return
    startup_deadline=time.monotonic()+220
    root=Path(value);config=json.loads((root/'app-runtime.json').read_text(encoding='utf-8-sig'))
    if config.get('engine')!='comfyui':return
    if config.get('backend')!='python':raise ValueError('HF custom-node startup requires backend python / engine comfyui')
    work=Path('/workspace/hf-app');user=work/'comfy-user';user.mkdir(parents=True,exist_ok=True)
    extra=work/'hf-extra-paths.yaml'
    wheels=sorted((root/'runtime_wheels').glob('*.whl'))
    packages=work/'python-deps'
    if wheels:
        stamp=json.dumps([(wheel.name,wheel.stat().st_size) for wheel in wheels])
        marker=packages/'.installed-wheels.json'
        if not marker.is_file() or marker.read_text()!=stamp:
            subprocess.run(['/opt/venv/bin/python','-m','pip','install','--disable-pip-version-check','--no-deps','--no-index','--no-compile','--target',str(packages),*[str(wheel) for wheel in wheels]],check=True,timeout=120)
            marker.write_text(stamp)
        sys.path.insert(0,str(packages))
        os.environ['PYTHONPATH']=str(packages)+os.pathsep+os.environ.get('PYTHONPATH','')
    # JSON is a valid YAML document and avoids path quoting ambiguities.
    paths=extra_paths(root)
    mappings={'hf_app':paths}
    if (root/'auxiliary'/'pose').is_dir():
        os.environ['AUX_ANNOTATOR_CKPTS_PATH']=str(root/'auxiliary'/'pose')
        os.environ['AUX_USE_SYMLINKS']='False'
        os.environ['AUX_TEMP_DIR']=str(work/'aux-temp')

    extra.write_text(json.dumps(mappings))
    entrypoints=[p for p in Path('/opt').glob('*/comfy_main_with_shared_models.py') if p.is_file()]
    if len(entrypoints)!=1:raise RuntimeError('Expected one installed runtime entrypoint')
    arguments=['/opt/venv/bin/python','-u',str(entrypoints[0]),
        '--listen','127.0.0.1','--port','8188','--disable-auto-launch','--disable-metadata',
        '--disable-api-nodes','--log-stdout','--enable-dynamic-vram','--reserve-vram','1',
        '--use-sage-attention','--async-offload','2','--preview-method','none',
        '--enable-triton-backend','--fast','fp16_accumulation',
        '--input-directory','/comfyui/input','--output-directory','/comfyui/output',
        '--temp-directory','/comfyui/temp','--user-directory',str(user),
        '--database-url','sqlite:////workspace/hf-app/comfy-user/comfyui.db','--extra-model-paths-config',str(extra)]
    PROCESS=subprocess.Popen(arguments,cwd='/comfyui')
    # Keep the enclosing v53 child ready/cleanup checks pointed at the real Comfy process.
    main=sys.modules.get('__main__')
    if main and hasattr(main,'COMFY'):main.COMFY=PROCESS
    deadline=min(startup_deadline,time.monotonic()+180)
    while time.monotonic()<deadline:
        if PROCESS.poll() is not None:raise RuntimeError('ComfyUI exited while loading HF runtime extensions')
        try:
            with urllib.request.urlopen('http://127.0.0.1:8188/system_stats',timeout=2) as response:
                if response.status==200:
                    print('HF_CATALOG_COMFY_READY',flush=True);return
        except Exception:pass
        time.sleep(.5)
    stop();raise TimeoutError('HF Comfy runtime did not become ready within180seconds')

atexit.register(stop)