import numpy as np
import requests
import cv2
import base64
import subprocess
import random
import time
import threading
import os

class IOPaintInpainting:
    """IOPaint 서버 연동 인페인팅 모델 (REST API /api/v1/inpaint 사용, 바이너리 PNG 반환)"""
    def __init__(self, server_url="http://localhost:8099"):
        self.api_url = f"{server_url}/api/v1/inpaint"
    def inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
        # 이미지를 base64로 인코딩
        _, img_encoded = cv2.imencode('.png', image)
        _, mask_encoded = cv2.imencode('.png', mask)
        img_b64 = base64.b64encode(img_encoded).decode('utf-8')
        mask_b64 = base64.b64encode(mask_encoded).decode('utf-8')
        payload = {
            "image": img_b64,
            "mask": mask_b64
        }
        response = requests.post(self.api_url, json=payload)
        if response.status_code != 200:
            print("IOPaint 서버 에러:", response.text)
            return None
        # 응답이 바이너리 PNG 이미지이므로 바로 디코딩
        nparr = np.frombuffer(response.content, np.uint8)
        result = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        return result

class IOPaintServerInstance:
    """IOPaint 서버 인스턴스 정보 및 상태 관리"""
    def __init__(self, port, process):
        self.port = port
        self.process = process
        self.busy = False
        self.last_used = time.time()
    def mark_busy(self):
        self.busy = True
        self.last_used = time.time()
    def mark_idle(self):
        self.busy = False
        self.last_used = time.time()
    def is_alive(self):
        return self.process.poll() is None

class IOPaintServerManager:
    """여러 IOPaint 서버 인스턴스를 관리하는 매니저"""
    def __init__(self, num_instances=1, port_range=(8099, 8199), iopaint_cmd=None, wait_ready=5):
        self.instances = []
        self.port_range = port_range
        self.iopaint_cmd = iopaint_cmd or ["python", "-m", "iopaint.server"]
        self.lock = threading.Lock()
        self._start_instances(num_instances, wait_ready)
    def _get_random_port(self):
        used_ports = {inst.port for inst in self.instances}
        candidates = [p for p in range(self.port_range[0], self.port_range[1]+1) if p not in used_ports]
        if not candidates:
            raise RuntimeError("사용 가능한 포트가 없습니다.")
        return random.choice(candidates)
    def _start_instances(self, num, wait_ready):
        for _ in range(num):
            port = self._get_random_port()
            cmd = self.iopaint_cmd + ["--port", str(port)]
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            instance = IOPaintServerInstance(port, proc)
            self.instances.append(instance)
            time.sleep(wait_ready)  # 서버가 준비될 때까지 대기
    def get_instance_info(self):
        """모든 인스턴스의 정보를 반환"""
        info = []
        for inst in self.instances:
            info.append({
                "port": inst.port,
                "busy": inst.busy,
                "alive": inst.is_alive(),
                "last_used": inst.last_used
            })
        return info
    def get_idle_instance(self):
        """놀고 있는(사용 가능한) 인스턴스 반환 (없으면 None)"""
        with self.lock:
            for inst in self.instances:
                if not inst.busy and inst.is_alive():
                    inst.mark_busy()
                    return inst
        return None
    def mark_instance_idle(self, port):
        """작업이 끝난 인스턴스를 idle로 표시"""
        for inst in self.instances:
            if inst.port == port:
                inst.mark_idle()
                break
    def shutdown_all(self):
        """모든 서버 인스턴스 종료"""
        for inst in self.instances:
            if inst.is_alive():
                inst.process.terminate()
        self.instances = []
    def add_instance(self, wait_ready=5):
        """새로운 인스턴스 추가"""
        port = self._get_random_port()
        cmd = self.iopaint_cmd + ["--port", str(port)]
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        instance = IOPaintServerInstance(port, proc)
        self.instances.append(instance)
        time.sleep(wait_ready)
        return instance