import subprocess
import logging
import socket
import random
import time
import os
import sys

class IOPaint_Server:
    def __init__(self, logger, base_dir):
        self.logger = logger
        self.base_dir = base_dir

    def find_free_port(self, port_range=(7020, 7030)):
        """포트 범위 내에서 사용 가능한 포트 반환"""
        for _ in range(20):
            port = random.randint(*port_range)
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                try:
                    s.bind(("0.0.0.0", port))
                    return port
                except OSError:
                    continue
        raise RuntimeError("사용 가능한 포트를 찾을 수 없습니다.")

    def start(self, wait_ready=30):
        """iopaint 서버를 실행하고, 정상 실행시 포트 반환"""
        port = self.find_free_port()
        model_dir = os.path.join(self.base_dir, 'iop', 'models')
        cmd = [sys.executable, '-m', 'iopaint', 'start', '--model=migan', '--device=cpu', '--port', str(port), '--model-dir', model_dir]
        self.logger.log(f"[IOPaint] 실행 환경 파이썬: {sys.executable}", level=logging.INFO)
        self.logger.log(f"[IOPaint] 실행 명령: {' '.join(cmd)}", level=logging.INFO)
        self.logger.log(f"[IOPaint] 모델 디렉토리: {model_dir}", level=logging.INFO)
        # pip list로 iopaint 설치여부 확인
        try:
            pip_list = subprocess.check_output([sys.executable, '-m', 'pip', 'list'], text=True)
            found = any('iopaint' in line for line in pip_list.splitlines())
            if found:
                self.logger.log("[IOPaint] iopaint 모듈이 현재 환경에 설치되어 있습니다.", level=logging.INFO)
            else:
                self.logger.log("[IOPaint] iopaint 모듈이 현재 환경에 설치되어 있지 않습니다!", level=logging.WARNING)
        except Exception as e:
            self.logger.log(f"[IOPaint] pip list 실행 실패: {e}", level=logging.ERROR, exc_info=True)

        # subprocess를 실시간 출력으로 실행
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1)
        self.logger.log(f"[IOPaint] 서버 준비 확인 시작 (최대 {wait_ready}초 대기)", level=logging.INFO)
        url = f"http://localhost:{port}/api/v1/server-config"
        start_time = time.time()
        stdout_lines = []
        stderr_lines = []
        import threading
        def read_stream(stream, lines, stream_name):
            for line in iter(stream.readline, ''):
                lines.append(line)
                print(f"[{stream_name}] {line}", end='')
            stream.close()
        t_out = threading.Thread(target=read_stream, args=(proc.stdout, stdout_lines, 'STDOUT'))
        t_err = threading.Thread(target=read_stream, args=(proc.stderr, stderr_lines, 'STDERR'))
        t_out.start()
        t_err.start()
        ready = False
        while time.time() - start_time < wait_ready:
            try:
                import requests
                r = requests.get(url, timeout=2)
                if r.status_code == 200:
                    self.logger.log(f"[IOPaint] 서버가 포트 {port}에서 준비됨.")
                    ready = True
                    break
            except Exception as e:
                time.sleep(0.5)
        t_out.join(timeout=2)
        t_err.join(timeout=2)
        if ready:
            return port
        # 실패 시 로그 및 예외
        self.logger.log(f"[IOPaint] 서버 실행 실패.\nstdout:\n{''.join(stdout_lines)}\nstderr:\n{''.join(stderr_lines)}", level=logging.ERROR, exc_info=True)
        print("[IOPaint] 서버 실행 실패. 전체 STDOUT:")
        print(''.join(stdout_lines))
        print("[IOPaint] 서버 실행 실패. 전체 STDERR:")
        print(''.join(stderr_lines))
        raise RuntimeError(f"IOPaint 서버가 {wait_ready}초 내에 준비되지 않았습니다.")

# if __name__ == "__main__":
#     logging.basicConfig(level=logging.INFO)
#     logger = logging.getLogger(__name__)
#     base_dir = os.path.dirname(os.path.abspath(__file__))
#     iop = IOPaint(logger, base_dir)
#     iop.start()