#!/usr/bin/env python3
"""
딸깍 개발실 — Real-ESRGAN 로컬 업스케일 서버 (의존성 0)

표준 라이브러리만 사용. 별도 pip install 필요 X.
실행 파일은 같은 폴더의 realesrgan/ 하위에 자동 배치되어 있어야 함 (setup.bat이 자동 처리).

사용법:
    python server.py              # 기본 포트 8765, localhost(127.0.0.1)만
    set PORT=8181 && python server.py   # 포트 변경 (Windows)
    HOST=0.0.0.0 python server.py       # 외부 접속 허용 (방화벽 주의)

엔드포인트:
    GET  /ping       — 헬스 체크
    GET  /health     — 헬스 체크 (별칭)
    POST /upscale    — 업스케일 요청 (JSON: {image, scale, model})

라이선스: 이 server.py는 GDI-Apps 프로젝트 라이선스를 따름.
번들된 realesrgan-ncnn-vulkan은 BSD-3-Clause (xinntao/Real-ESRGAN-ncnn-vulkan).
"""

import os
import sys
import json
import base64
import tempfile
import subprocess
import logging
import platform
import socket
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

# ─── 설정 ─────────────────────────────────────
HERE = os.path.dirname(os.path.abspath(__file__))
EXE_NAME = "realesrgan-ncnn-vulkan.exe" if platform.system() == "Windows" else "realesrgan-ncnn-vulkan"
ESRGAN_BIN = os.environ.get("ESRGAN_BIN", os.path.join(HERE, "realesrgan", EXE_NAME))
ESRGAN_MODEL_DIR = os.environ.get("ESRGAN_MODEL_DIR", os.path.join(HERE, "realesrgan", "models"))
DEFAULT_MODEL = os.environ.get("ESRGAN_DEFAULT_MODEL", "realesrgan-x4plus")
# v2.0.1: HOST 기본값 127.0.0.1로 변경 (Windows에서 0.0.0.0 바인딩이 권한 거부되는 환경 회피)
HOST = os.environ.get("HOST", "127.0.0.1")
# v2.0.1: PORT_CANDIDATES — 첫 포트가 막히면 자동으로 다음 포트 시도 (Hyper-V/Docker 예약 회피)
PORT_PRIMARY = int(os.environ.get("PORT", 8765))
PORT_CANDIDATES = [PORT_PRIMARY, 8181, 8282, 8383, 9765, 17654]

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
log = logging.getLogger("esrgan-server")


# ─── 핵심 로직 ─────────────────────────────────
def _status():
    return {
        "ok": True,
        "status": "ok",
        "service": "딸깍 개발실 - Real-ESRGAN Local Server",
        "version": "2.0.1",
        "esrgan_bin_found": os.path.exists(ESRGAN_BIN),
        "esrgan_bin": ESRGAN_BIN,
        "model_dir": ESRGAN_MODEL_DIR,
        "default_model": DEFAULT_MODEL,
    }


def run_upscale(image_b64, scale=4, model=None):
    """ncnn-vulkan 실행파일을 subprocess로 호출. 결과를 base64로 반환."""
    if not os.path.exists(ESRGAN_BIN):
        raise FileNotFoundError(f"Real-ESRGAN 실행파일을 찾을 수 없습니다: {ESRGAN_BIN}")

    # dataURL prefix 제거
    if "," in image_b64:
        image_b64 = image_b64.split(",", 1)[1]

    try:
        img_bytes = base64.b64decode(image_b64)
    except Exception as e:
        raise ValueError(f"base64 디코딩 실패: {e}")

    model = model or DEFAULT_MODEL
    scale = max(2, min(4, int(scale)))  # 안전 범위

    with tempfile.TemporaryDirectory() as tmpdir:
        in_path = os.path.join(tmpdir, "input.png")
        out_path = os.path.join(tmpdir, "output.png")

        with open(in_path, "wb") as f:
            f.write(img_bytes)

        cmd = [
            ESRGAN_BIN,
            "-i", in_path,
            "-o", out_path,
            "-s", str(scale),
            "-n", model,
            "-m", ESRGAN_MODEL_DIR,
        ]
        log.info("Running: %s", " ".join(cmd))

        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        except subprocess.TimeoutExpired:
            raise TimeoutError("Real-ESRGAN 처리 시간 초과 (5분)")
        except FileNotFoundError as e:
            raise FileNotFoundError(f"실행 실패 — 파일은 있지만 실행 권한 또는 의존성 문제일 수 있음: {e}")

        if result.returncode != 0:
            err = (result.stderr or result.stdout or "")[:500]
            raise RuntimeError(f"Real-ESRGAN 오류 (code {result.returncode}): {err}")

        with open(out_path, "rb") as f:
            return base64.b64encode(f.read()).decode("ascii")


# ─── HTTP 핸들러 ─────────────────────────────────
class Handler(BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):
        # 기본 stderr 로그 대신 우리 logger 사용
        log.info("%s - %s", self.address_string(), fmt % args)

    def _send_json(self, payload, status=200):
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.end_headers()
        self.wfile.write(body)

    def do_OPTIONS(self):
        self.send_response(204)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.send_header("Access-Control-Max-Age", "86400")
        self.end_headers()

    def do_GET(self):
        if self.path in ("/ping", "/health", "/"):
            self._send_json(_status())
        else:
            self._send_json({"ok": False, "error": f"Unknown endpoint: {self.path}"}, 404)

    def do_POST(self):
        if self.path != "/upscale":
            self._send_json({"ok": False, "error": f"Unknown endpoint: {self.path}"}, 404)
            return

        length = int(self.headers.get("Content-Length", 0))
        if length <= 0 or length > 50 * 1024 * 1024:  # 50MB 한도
            self._send_json({"ok": False, "error": "잘못된 Content-Length (0 또는 50MB 초과)"}, 400)
            return

        raw = self.rfile.read(length)
        try:
            data = json.loads(raw.decode("utf-8"))
        except Exception as e:
            self._send_json({"ok": False, "error": f"JSON 파싱 실패: {e}"}, 400)
            return

        image = data.get("image")
        if not image:
            self._send_json({"ok": False, "error": "image 필드가 필요합니다"}, 400)
            return

        scale = data.get("scale", 4)
        model = data.get("model")

        try:
            out_b64 = run_upscale(image, scale=scale, model=model)
            self._send_json({
                "ok": True,
                "image": out_b64,
                "model": model or DEFAULT_MODEL,
                "scale": scale,
            })
        except FileNotFoundError as e:
            self._send_json({"ok": False, "error": str(e)}, 500)
        except ValueError as e:
            self._send_json({"ok": False, "error": str(e)}, 400)
        except TimeoutError as e:
            self._send_json({"ok": False, "error": str(e)}, 504)
        except Exception as e:
            log.exception("업스케일 처리 중 예외")
            self._send_json({"ok": False, "error": f"서버 오류: {e}"}, 500)


# ─── 메인 ─────────────────────────────────
def _try_bind(host, port):
    """포트 바인딩 시도. 성공하면 ThreadingHTTPServer, 실패하면 None."""
    try:
        return ThreadingHTTPServer((host, port), Handler)
    except OSError as e:
        return e


def main():
    bin_ok = os.path.exists(ESRGAN_BIN)
    print("=" * 60)
    print(" 딸깍 개발실 — Real-ESRGAN 로컬 업스케일 서버")
    print("=" * 60)
    print(f" 실행 파일  : {ESRGAN_BIN}")
    print(f"              {'[OK] 발견됨' if bin_ok else '[ERROR] 못 찾음 - setup.bat 먼저 실행하세요'}")
    print(f" 모델 폴더  : {ESRGAN_MODEL_DIR}")
    print(f" 기본 모델  : {DEFAULT_MODEL}")
    print()
    print(" 종료 방법  : 이 창을 닫거나 Ctrl+C")
    print("=" * 60)
    print()

    if not bin_ok:
        print("[ERROR] Real-ESRGAN 실행 파일이 없습니다.")
        print("       ZIP을 다시 풀어주세요.")
        print()

    # v2.0.1: 포트 fallback — 첫 포트가 막혔으면 (Windows Hyper-V/Docker 예약 등) 다음 후보 시도
    httpd = None
    used_port = None
    last_err = None
    for port in PORT_CANDIDATES:
        result = _try_bind(HOST, port)
        if isinstance(result, ThreadingHTTPServer):
            httpd = result
            used_port = port
            break
        last_err = (port, result)
        log.warning(f"포트 {port} 바인딩 실패: {result}")

    if httpd is None:
        print()
        print("=" * 60)
        print(" [ERROR] 사용할 수 있는 포트를 찾을 수 없습니다.")
        print("=" * 60)
        if last_err and hasattr(last_err[1], 'winerror') and last_err[1].winerror == 10013:
            print(" 원인: WinError 10013 — Windows의 Hyper-V / Docker / WSL2가")
            print("       부팅 시 일부 포트 범위를 예약해서 막은 것으로 보입니다.")
            print()
            print(" 해결 방법:")
            print("  1) cmd 열고 아래 명령으로 예약된 포트 확인:")
            print("       netsh interface ipv4 show excludedportrange protocol=tcp")
            print()
            print("  2) 예약되지 않은 포트를 직접 지정해서 실행:")
            print("       set PORT=7777 && python\\python.exe server.py")
            print("     (그 다음 에디터에서 서버 URL을 http://localhost:7777로 변경)")
            print()
            print("  3) 또는 관리자 권한 cmd에서 winnat 재시작:")
            print("       net stop winnat")
            print("       net start winnat")
        else:
            print(f" 마지막 시도 포트 {last_err[0]}: {last_err[1]}")
            print()
            print(" 다른 포트로 시도해보세요:")
            print("   set PORT=7777 && python\\python.exe server.py")
        print("=" * 60)
        sys.exit(1)

    print(f" [OK] 서버 시작됨")
    print(f" 서버 주소  : http://localhost:{used_port}")
    if used_port != PORT_PRIMARY:
        print(f"              (기본 포트 {PORT_PRIMARY}가 막혀서 {used_port}로 fallback)")
        print(f"              에디터에서 서버 URL을 http://localhost:{used_port} 로 변경하세요!")
    print()

    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print("\n서버 종료.")
    finally:
        httpd.server_close()


if __name__ == "__main__":
    main()
