import subprocess
import sys
import time
import os

def run_webui():
    user_args = sys.argv[1:]
    python_exe = "/content/venv/GUTRIS1/bin/python3.10"
    command = [python_exe, "launch.py"] + user_args

    attempt = 1
    while True:
        print(f"\n🚀 [{attempt}회차] WebUI 실행 시도...")
        
        # stdout=None으로 설정하면 시스템이 직접 로그를 찍으므로 속도 저하가 0입니다.
        process = subprocess.Popen(
            command,
            env=os.environ.copy(),
            cwd="/content/stable-diffusion-webui-reForge"
        )

        # 프로세스가 끝날 때까지 기다립니다.
        return_code = process.wait()

        # 만약 에러 코드를 남기고 죽었다면 (보통 1) 자동으로 재시작
        if return_code != 0:
            print(f"\n❌ 에러 발생 (Exit Code: {return_code}). 2초 후 자동 재시작합니다...")
            attempt += 1
            time.sleep(2)
            continue
        else:
            # 정상 종료(0)인 경우 루프 탈출
            print("\n✅ WebUI 종료.")
            break

if __name__ == "__main__":
    run_webui()