import os
import subprocess
import time
import atexit
import sys
import re
import argparse
import shlex
import io
import select

# --- 설정 변수 ---
COMFYUI_PORT = 8188

# 🚨 사용자 지정: 모든 .py 파일을 저장할 절대 경로
# 이 경로에 gradio-tunnel.py가 저장되어야 합니다. (Colab 셀 실행 전에 /content/py 폴더가 존재해야 합니다.)
PYTHON_SCRIPT_DIR = "/content/py"

# 터널 스크립트의 최종 절대 경로
TUNNEL_SCRIPT_PATH = os.path.join(PYTHON_SCRIPT_DIR, "gradio-tunnel.py") 
# 예: "/content/py/gradio-tunnel.py"

# Gradio Tunnel Script 내용 (터널링 기능을 수행하는 Python 코드)
GRADIO_TUNNEL_SCRIPT_CONTENT = r"""
import atexit
import os
import platform
import re
import stat
import subprocess
import sys
import time
import secrets
from pathlib import Path
from typing import List, Optional

import requests

VERSION = "0.2"
CURRENT_TUNNELS: List["Tunnel"] = []

machine = platform.machine()
if machine == "x86_64":
    machine = "amd64"

BINARY_REMOTE_NAME = f"frpc_{platform.system().lower()}_{machine.lower()}"
EXTENSION = ".exe" if os.name == "nt" else ""
BINARY_URL = f"https://cdn-media.huggingface.co/frpc-gradio-{VERSION}/{BINARY_REMOTE_NAME}{EXTENSION}"

BINARY_FILENAME = f"{BINARY_REMOTE_NAME}_v{VERSION}"
BINARY_FOLDER = Path(__file__).parent.absolute()
BINARY_PATH = f"{BINARY_FOLDER / BINARY_FILENAME}"

TUNNEL_TIMEOUT_SECONDS = 30
TUNNEL_ERROR_MESSAGE = (
    "Could not create share URL. "
    "Please check the appended log from frpc for more information:"
)

GRADIO_API_SERVER = "https://api.gradio.app/v2/tunnel-request"
GRADIO_SHARE_SERVER_ADDRESS = None


class Tunnel:
    def __init__(self, remote_host, remote_port, local_host, local_port, share_token):
        self.proc = None
        self.url = None
        self.remote_host = remote_host
        self.remote_port = remote_port
        self.local_host = local_host
        self.local_port = local_port
        self.share_token = share_token

    @staticmethod
    def download_binary():
        if not Path(BINARY_PATH).exists():
            resp = requests.get(BINARY_URL)

            if resp.status_code == 403:
                raise OSError(
                    f"Cannot set up a share link as this platform is incompatible. Please "
                    f"create a GitHub issue with information about your platform: {platform.uname()}"
                )

            resp.raise_for_status()

            # Save file data to local copy
            with open(BINARY_PATH, "wb") as file:
                file.write(resp.content)
            st = os.stat(BINARY_PATH)
            os.chmod(BINARY_PATH, st.st_mode | stat.S_IEXEC)

    def start_tunnel(self) -> str:
        self.download_binary()
        self.url = self._start_tunnel(BINARY_PATH)
        return self.url

    def kill(self):
        if self.proc is not None:
            print(f"터널 종료 중: {self.local_host}:{self.local_port} <> {self.url}")
            self.proc.terminate()
            self.proc = None

    def _start_tunnel(self, binary: str) -> str:
        CURRENT_TUNNELS.append(self)
        command = [
            binary,
            "http",
            "-n",
            self.share_token,
            "-l",
            str(self.local_port),
            "-i",
            self.local_host,
            "--uc",
            "--sd",
            "random",
            "--ue",
            "--server_addr",
            f"{self.remote_host}:{self.remote_port}",
            "--disable_log_color",
        ]
        self.proc = subprocess.Popen(
            command, stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )
        atexit.register(self.kill)
        return self._read_url_from_tunnel_stream()

    def _read_url_from_tunnel_stream(self) -> str:
        start_timestamp = time.time()

        log = []
        url = ""

        def _raise_tunnel_error():
            log_text = "\\n".join(log) 
            print(log_text, file=sys.stderr)
            raise ValueError(f"{TUNNEL_ERROR_MESSAGE}\\n{log_text}")

        while url == "":
            if time.time() - start_timestamp >= TUNNEL_TIMEOUT_SECONDS:
                _raise_tunnel_error()

            assert self.proc is not None
            if self.proc.stdout is None:
                continue

            line = self.proc.stdout.readline()
            try:
                line = line.decode("utf-8")
            except UnicodeDecodeError:
                continue
            
            if line == "":
                time.sleep(0.01)
                continue

            log.append(line.strip())

            if "start proxy success" in line:
                result = re.search("start proxy success: (.+)\\n", line)
                if result is None:
                    _raise_tunnel_error()
                else:
                    url = result.group(1)
            elif "login to server failed" in line:
                _raise_tunnel_error()
        
        if self.proc.poll() is not None and url == "":
            _raise_tunnel_error()

        return url


def setup_tunnel(
    local_host: str,
    local_port: int,
    share_token: str,
    share_server_address: Optional[str],
) -> str:
    share_server_address = (
        GRADIO_SHARE_SERVER_ADDRESS
        if share_server_address is None
        else share_server_address
    )
    if share_server_address is None:
        response = requests.get(GRADIO_API_SERVER)
        if not (response and response.status_code == 200):
            raise RuntimeError("Could not get share link from Gradio API Server.")
        payload = response.json()[0]
        remote_host, remote_port = payload["host"], int(payload["port"])
    else:
        remote_host, remote_port = share_server_address.split(":")
        remote_port = int(remote_port)
    try:
        tunnel = Tunnel(remote_host, remote_port, local_host, local_port, share_token)
        address = tunnel.start_tunnel()
        return address
    except Exception as e:
        raise RuntimeError(str(e)) from e

# 이 스크립트는 포트 번호를 인자로 받아 Gradio 터널링을 시작합니다.
if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("사용법: python gradio-tunnel.py <port_number>", file=sys.stderr)
        sys.exit(1)
        
    try:
        port = int(sys.argv[1])
    except ValueError:
        print("오류: 포트 번호는 정수여야 합니다.", file=sys.stderr)
        sys.exit(1)

    # requests 임포트 확인 (Colab 환경에서 필요)
    try:
        import requests
    except ImportError:
        # 이 코드는 runcomfy.py에서 실행하는 subprocess이므로, 
        # runcomfy.py에서 requests 설치를 보장해야 함.
        pass 

    try:
        address = setup_tunnel(
            "127.0.0.1",
            port,
            secrets.token_urlsafe(32),
            None,
        )
        print(address, flush=True)
        # 터널 프로세스가 종료되지 않도록 대기
        time.sleep(3600 * 24 * 3)
    except Exception as e:
        print(f"터널 실행 중 치명적인 오류 발생: {e}", file=sys.stderr)
        sys.exit(1)

"""

# --- 1. main 함수 ---
def main():
    # requests 라이브러리 설치를 보장합니다.
    try:
        import requests
    except ImportError:
        print("필수 라이브러리 'requests'가 설치되어 있지 않습니다. 설치를 시도합니다.")
        try:
            # 외부 라이브러리 설치를 시도
            subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"])
        except Exception as e:
            print(f"경고: 'requests' 라이브러리 설치 실패. 터널링 오류가 발생할 수 있습니다. 오류: {e}")

    # runcomfy.py에 전달된 모든 인자(스크립트 이름 제외)를 ComfyUI에 전달할 인자로 사용
    comfyui_args = sys.argv[1:]

    # --- 필수 경로 설정 및 이동 ---
    # WS 변수는 이 스크립트를 실행하는 외부 Colab 셀에서 결정되거나, 
    # 혹은 이 스크립트가 DRIVE 모드(/content/drive/MyDrive/ONECLICK) 전용으로 실행된다고 가정합니다.
    comfyui_dir = "/content/ComfyUI" 
    
    print(f"ComfyUI 디렉토리로 이동 중: {comfyui_dir}")
    try:
        os.chdir(comfyui_dir)
    except FileNotFoundError:
        print(f"오류: 디렉토리가 없습니다: {comfyui_dir}. 경로를 확인해 주세요.")
        sys.exit(1) # 디렉토리가 없으면 종료

    # 🚨 수정: 터널 스크립트 파일을 지정된 경로에 저장 (폴더 생성 없음)
    print(f"터널 스크립트 저장 중: {TUNNEL_SCRIPT_PATH}")
    try:
        # TUNNEL_SCRIPT_PATH에 저장 (절대 경로 사용)
        with open(TUNNEL_SCRIPT_PATH, "w") as f:
            f.write(GRADIO_TUNNEL_SCRIPT_CONTENT)
    except Exception as e:
        print(f"오류: 터널 스크립트 파일 저장 실패. {PYTHON_SCRIPT_DIR} 경로가 존재하는지 확인해 주세요. 오류: {e}")
        sys.exit(1)


    # --- 2. ComfyUI 서버 실행 (백그라운드) ---
    print("\nComfyUI 서버를 백그라운드로 시작합니다...")
    comfyui_base_args = [
        "./main.py", 
        "--listen", "127.0.0.1", 
        "--port", str(COMFYUI_PORT),
    ]
    comfyui_command = ["python"] + comfyui_base_args + comfyui_args
    
    comfyui_proc = subprocess.Popen(
        comfyui_command,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT, 
        text=True 
    )
    atexit.register(comfyui_proc.terminate)

    # --- 3. Gradio 터널링 스크립트 실행 (백그라운드) ---
    print(f"서버가 포트 {COMFYUI_PORT}를 열 때까지 15초 대기합니다...")
    time.sleep(15)

    print(f"\nGradio 터널 클라이언트 시작 중 (포트 {COMFYUI_PORT})...")
    # 🚨 수정: TUNNEL_SCRIPT_PATH 변수(절대 경로) 사용
    tunnel_command = ["python", TUNNEL_SCRIPT_PATH, str(COMFYUI_PORT)]
    
    tunnel_proc = subprocess.Popen(
        tunnel_command,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT, 
        text=True
    )
    atexit.register(tunnel_proc.terminate)

    # --- 4. 터널 주소 출력 및 프로세스 유지 ---
    share_url = None
    start_time = time.time()
    print("\n[터널 출력 - Gradio 로그]")

    # 터널 URL을 찾는 루프
    while time.time() - start_time < 60 and tunnel_proc.poll() is None:
        line = tunnel_proc.stdout.readline()
        
        if line:
            line_stripped = line.strip()
            print(f"[TUNNEL] {line_stripped}")
            
            # URL 패턴 찾기
            if line_stripped.startswith(("https://", "http://")):
                share_url = line_stripped
                break
                
        time.sleep(0.01)

    if share_url:
        print("\n" + "="*50)
        print(f"🎉 **ComfyUI 공유 URL이 생성되었습니다:** {share_url}")
        print("이 주소로 접속하세요.")
        print("="*50)

    else:
        # --- 터널링 실패 시 오류 진단 로직 ---
        print("\n" + "="*50)
        print("🚨 **터널링 클라이언트 실행에 실패했습니다!**")
        print("아래 로그에서 오류 원인을 확인하세요.")
        print("="*50)
        
        # 1. 터널 프로세스 종료 시 남은 로그 확인 및 출력
        print("\n--- 터널 클라이언트 최종 로그 ---")
        # 터널 stdout의 처음으로 이동 (이미 읽은 내용 포함)
        try:
            tunnel_logs = tunnel_proc.stdout.read()
            if tunnel_logs:
                 print(tunnel_logs.strip())
            else:
                print("터널 로그가 발견되지 않았습니다. 즉시 종료된 것으로 보입니다.")
        except Exception as e:
            print(f"터널 로그를 읽는 중 예외 발생: {e}")

        # 2. ComfyUI 프로세스 로그 확인
        print("\n--- ComfyUI 초기 로그 확인 (ComfyUI가 포트를 열었는지 확인) ---")
        try:
            if comfyui_proc.stdout is not None:
                # select를 사용하여 비차단 방식으로 읽을 수 있도록 시도
                ready_to_read, _, _ = select.select([comfyui_proc.stdout], [], [], 0.5)
                if ready_to_read:
                    comfy_logs = comfyui_proc.stdout.read()
                    print(comfy_logs.strip() if comfy_logs else "ComfyUI 로그 버퍼에 내용 없음.")
                else:
                    print("ComfyUI 로그 버퍼에 내용 없음 (500ms 대기 후).")
            else:
                print("ComfyUI stdout이 None입니다.")
        except Exception as e:
            print(f"ComfyUI 로그 읽기 중 예외 발생: {e}")

        # 모든 프로세스 종료 후 오류 발생
        if comfyui_proc.poll() is None:
            comfyui_proc.terminate()
        if tunnel_proc.poll() is None:
            tunnel_proc.terminate()

        raise RuntimeError("Gradio 터널링에 실패했습니다. 위의 상세 로그를 확인해 주십시오.")


    # 프로세스 유지를 위한 메인 루프 (충돌 로그 출력 강화)
    try:
        print("\n**프로세스를 유지합니다. Colab 세션이 연결된 동안 터널이 유지됩니다.**")
        print("중지하려면 이 셀의 실행을 멈추세요.")
        while True:
            # ComfyUI 로그를 실시간으로 모두 출력
            comfy_line = comfyui_proc.stdout.readline()
            if comfy_line:
                print(comfy_line.strip())

            if comfyui_proc.poll() is not None:
                print("\n[오류] ComfyUI 프로세스가 예기치 않게 중지되었습니다. 세션을 종료합니다.")
                # 종료 시 남아있는 모든 로그를 출력하여 원인 파악
                print("\n--- ComfyUI 충돌 로그 (원인 분석용) ---")
                print(comfyui_proc.stdout.read())
                break
                
            if tunnel_proc.poll() is not None:
                print("\n[오류] 터널 프로세스가 예기치 않게 중지되었습니다. 세션을 종료합니다.")
                # 터널 프로세스 종료 시 남은 로그 확인 및 출력
                print("\n--- 터널 충돌 로그 (원인 분석용) ---")
                print(tunnel_proc.stdout.read())
                break
                
            time.sleep(0.01)

    except KeyboardInterrupt:
        print("\n사용자가 프로세스를 중지했습니다. (KeyboardInterrupt)")

    finally:
        print("\n백그라운드 프로세스 종료를 시도합니다...")
        
        if comfyui_proc.poll() is None:
            comfyui_proc.terminate()
        if tunnel_proc.poll() is None:
            tunnel_proc.terminate()
        
        print("백그라운드 프로세스가 종료되었습니다.")

if __name__ == '__main__':
    main()