IMG_Worker/modules/tray_app.py

import os
import sys
import threading
import time
from typing import Optional
import json

import pystray
from PIL import Image, ImageDraw
import requests

from loggerModule import Logger

# 상위 디렉토리를 sys.path에 추가하여 updater 모듈을 찾을 수 있게 함
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

try:
    from updater.__version__ import __version__ as app_version
except ImportError:
    app_version = "unknown"


class TrayController:
    def __init__(self):
        logs_root = os.path.join(os.environ.get("PROGRAMDATA", r"C:\\ProgramData"), "ImgWorker")
        os.makedirs(logs_root, exist_ok=True)
        self.logger = Logger(log_file=os.path.join(logs_root, "tray.log"), logger_name="img_worker_tray")
        self.api = _TrayAPI(self.logger)
        self.icon: Optional[pystray.Icon] = None
        self._stop_event = threading.Event()
        self._last_ready: bool = False

    def _create_image(self, color=(0, 120, 215)):
        img = Image.new('RGB', (64, 64), color)
        draw = ImageDraw.Draw(img)
        draw.rectangle([8, 8, 56, 56], outline=(255, 255, 255), width=3)
        draw.ellipse([22, 22, 42, 42], fill=(255, 255, 255))
        return img

    def _set_icon_color(self, state: str):
        # state: 'ready' -> green, 'stopped' -> red, 'active' -> blink green/yellow
        green = (0, 160, 0)
        red = (180, 0, 0)
        yellow = (200, 160, 0)
        color = red
        if state == 'ready':
            color = green
        elif state == 'active':
            # 깜박임은 폴링 스레드에서 교차로 변경
            color = green if int(time.time()) % 2 == 0 else yellow
        if self.icon:
            self.icon.icon = self._create_image(color)

    def _update_title(self):
        try:
            st = self.api.worker_status()
            ready = bool(st.get("ready", False))
            active = bool(st.get("active", False))
            avg_sec = st.get("avg_sec_per_image")
            provider = (st.get("provider") or "").upper()
            pid = st.get("pid")
            self._last_ready = ready
            status_label = 'ACTIVE' if active else ('READY' if ready else 'STOP')
            perf = f" - avg. {avg_sec:.1f}(sec/장)" if isinstance(avg_sec, (int, float)) else ''
            prov = f" [{provider}]" if provider else ''
            title = f"ImgWorker v{app_version} [{status_label}]{perf}{prov}"
            if self.icon:
                self.icon.title = title
                # 메뉴를 상태에 맞게 재구성
                self.icon.menu = self._build_menu()
                try:
                    self.icon.update_menu()
                except Exception:
                    pass
                self._set_icon_color('active' if active else ('ready' if ready else 'stopped'))
        except Exception:
            if self.icon:
                self.icon.title = "ImgWorker (offline)"

    def _build_menu(self) -> pystray.Menu:
        status_text = f"워커 상태: {'실행 중' if self._last_ready else '중지'}"
        # 상태 라벨(비활성), 시작/중지 버튼은 상태에 따라 enable
        return pystray.Menu(
            pystray.MenuItem(status_text, lambda: None, enabled=False),
            pystray.MenuItem("워커 시작", self._action_start, enabled=(not self._last_ready)),
            pystray.MenuItem("워커 중지", self._action_stop, enabled=self._last_ready),
            pystray.MenuItem("로그 열기", self._action_open_log),
            pystray.Menu.SEPARATOR,
            pystray.MenuItem("상태 새로고침", self._action_info),
            pystray.MenuItem("서버 종료", self._action_exit),
        )

    def _action_start(self, icon, item):
        self.logger.info("Tray: start worker")
        self.api.worker_start()
        time.sleep(0.3)
        self._update_title()

    def _action_stop(self, icon, item):
        self.logger.info("Tray: stop worker")
        self.api.worker_stop()
        time.sleep(0.3)
        self._update_title()

    def _action_info(self, icon, item):
        self.logger.info("Tray: info refresh")
        self._update_title()

    def _action_exit(self, icon, item):
        self.logger.info("Tray: exit clicked")
        try:
            self.api.shutdown_server()
        except Exception:
            pass
        self._stop_event.set()
        if self.icon:
            self.icon.stop()

    def _action_open_log(self, icon, item):
        try:
            logs_root = os.path.join(os.environ.get("PROGRAMDATA", r"C:\\ProgramData"), "ImgWorker")
            log_path = os.path.join(logs_root, "logs", "api_server.log")
            # 기본 편집기로 열기
            if os.name == "nt":
                os.startfile(log_path)  # type: ignore[attr-defined]
            else:
                import subprocess
                subprocess.Popen(["xdg-open", log_path])
        except Exception as e:
            self.logger.error(f"로그 열기 실패: {e}")

    def run(self):
        self.icon = pystray.Icon("imgworker_tray", self._create_image(), "ImgWorker", self._build_menu())

        # 상태 폴링 쓰레드로 제목/메뉴 갱신
        def _poll():
            while not self._stop_event.is_set():
                try:
                    self._update_title()
                except Exception:
                    pass
                time.sleep(2.0)

        t = threading.Thread(target=_poll, daemon=True)
        t.start()
        # 최초 상태 반영
        self._update_title()
        self.icon.run()


def main():
    TrayController().run()


if __name__ == "__main__":
    main()


def _read_server_info() -> Optional[dict]:
    try:
        program_data = os.environ.get("PROGRAMDATA", r"C:\\ProgramData")
        info_path = os.path.join(program_data, "ImgWorker", "server.json")
        if os.path.isfile(info_path):
            with open(info_path, "r", encoding="utf-8") as f:
                return json.load(f)
    except Exception:
        return None
    return None


class _TrayAPI:
    def __init__(self, logger: Logger):
        self.logger = logger
        base = os.environ.get("IMGWK_API_BASE", "")
        if not base:
            info = _read_server_info() or {}
            base = info.get("base") or f"http://{info.get('host','127.0.0.1')}:{info.get('port',8009)}"
        if not base:
            base = "http://127.0.0.1:8009"
        self.base = base.rstrip("/")

    def worker_status(self) -> dict:
        try:
            r = requests.get(f"{self.base}/v1/worker/status", timeout=5)
            r.raise_for_status()
            return r.json()
        except Exception as e:
            return {"ready": False, "error": str(e)}

    def worker_start(self) -> dict:
        try:
            r = requests.post(f"{self.base}/v1/worker/start", timeout=10)
            r.raise_for_status()
            return r.json()
        except Exception as e:
            return {"ok": False, "error": str(e)}

    def worker_stop(self) -> dict:
        try:
            r = requests.post(f"{self.base}/v1/worker/stop", timeout=10)
            r.raise_for_status()
            return r.json()
        except Exception as e:
            return {"ok": False, "error": str(e)}

    def shutdown_server(self) -> dict:
        try:
            r = requests.post(f"{self.base}/v1/server/shutdown", timeout=5)
            r.raise_for_status()
            return r.json()
        except Exception as e:
            return {"ok": False, "error": str(e)}