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 실행 시도...")
        
        process = subprocess.Popen(
            command,
            env=os.environ.copy(),
            cwd="/content/stable-diffusion-webui-reForge"
        )

        try:
            # 프로세스가 끝날 때까지 대기
            return_code = process.wait()
            
            if return_code != 0:
                print(f"\n❌ 에러 발생 (Exit Code: {return_code}). 2초 후 자동 재시작...")
                attempt += 1
                time.sleep(2)
                continue
            else:
                print("\n✅ WebUI 정상 종료.")
                break

        except KeyboardInterrupt:
            # 사용자가 Ctrl+C를 눌렀을 때
            print("\n\n👋 사용자가 중단 신호를 보냈습니다. 안전하게 종료합니다...")
            process.terminate() # 자식 프로세스(WebUI)도 같이 죽여줌
            try:
                process.wait(timeout=5)
            except:
                process.kill()
            print("✅ WebUI 종료 완료.")
            break # 루프 탈출

if __name__ == "__main__":
    run_webui()