Blame for runner.py

"""
Headless runner for student Pygame / OpenCV projects.

Connects back to the server over a private Unix socket, monkeypatches
display/input APIs, and streams JPEG frames + logs to the browser.
"""
from __future__ import annotations

import argparse
import io
import json
import os
import queue
import socket
import sys
import tempfile
import threading
import time

import numpy as np

# Parse CLI before importing pygame/opencv (env must be set first)
parser = argparse.ArgumentParser()
parser.add_argument("--project_path", required=True, help="Path to the student's main.py")
parser.add_argument("--socket-path", required=True, help="Private Unix socket used for IPC")
parser.add_argument(
    "--models_dir",
    default=None,
    help="Shared MediaPipe / model assets directory",
)
args = parser.parse_args()

os.environ["SDL_VIDEODRIVER"] = "dummy"
os.environ["SDL_AUDIODRIVER"] = "dummy"

# ----------------- IPC -----------------
send_queue: queue.Queue = queue.Queue()
cv2_keys: queue.Queue = queue.Queue()
stdin_queue: queue.Queue = queue.Queue()

# Adaptive frame pacing: start at ~30 FPS, back off under pressure
_min_frame_interval = 0.033
_max_frame_interval = 0.08
_current_frame_interval = 0.033
JPEG_QUALITY = 55

class KeyStateWrapper:
    def __init__(self):
        self._keys = {}

    def __getitem__(self, key):
        return self._keys.get(key, False)

    def __setitem__(self, key, value):
        self._keys[key] = bool(value)

    def __len__(self):
        return 512


shared_keys = KeyStateWrapper()
shared_mouse_pos = (0, 0)
shared_mouse_buttons = (False, False, False)
shared_webcam_frame = None

last_pygame_frame_time = 0.0
last_cv2_frame_time = 0.0

import pygame  # noqa: E402  — after SDL env

KEY_MAP = {
    "ArrowLeft": pygame.K_LEFT,
    "ArrowRight": pygame.K_RIGHT,
    "ArrowUp": pygame.K_UP,
    "ArrowDown": pygame.K_DOWN,
    " ": pygame.K_SPACE,
    "Spacebar": pygame.K_SPACE,
    "Enter": pygame.K_RETURN,
    "Escape": pygame.K_ESCAPE,
    "Backspace": pygame.K_BACKSPACE,
    "Tab": pygame.K_TAB,
    "Shift": pygame.K_LSHIFT,
    "Control": pygame.K_LCTRL,
    "Alt": pygame.K_LALT,
    "a": pygame.K_a, "b": pygame.K_b, "c": pygame.K_c, "d": pygame.K_d, "e": pygame.K_e,
    "f": pygame.K_f, "g": pygame.K_g, "h": pygame.K_h, "i": pygame.K_i, "j": pygame.K_j,
    "k": pygame.K_k, "l": pygame.K_l, "m": pygame.K_m, "n": pygame.K_n, "o": pygame.K_o,
    "p": pygame.K_p, "q": pygame.K_q, "r": pygame.K_r, "s": pygame.K_s, "t": pygame.K_t,
    "u": pygame.K_u, "v": pygame.K_v, "w": pygame.K_w, "x": pygame.K_x, "y": pygame.K_y,
    "z": pygame.K_z,
    "0": pygame.K_0, "1": pygame.K_1, "2": pygame.K_2, "3": pygame.K_3, "4": pygame.K_4,
    "5": pygame.K_5, "6": pygame.K_6, "7": pygame.K_7, "8": pygame.K_8, "9": pygame.K_9,
}


def _adjust_frame_interval(queue_depth: int) -> None:
    global _current_frame_interval
    if queue_depth > 2:
        _current_frame_interval = min(_max_frame_interval, _current_frame_interval * 1.15)
    elif queue_depth == 0:
        _current_frame_interval = max(_min_frame_interval, _current_frame_interval * 0.95)


# ----------------- private IPC threads -----------------
def reader_thread_func(sock):
    global shared_webcam_frame, shared_keys, shared_mouse_pos, shared_mouse_buttons
    print("[Runner-IPC] Reader thread started.")
    try:
        while True:
            type_bytes = sock.recv(4)
            if not type_bytes or len(type_bytes) < 4:
                break
            msg_type = type_bytes.decode("utf-8", errors="ignore")

            len_bytes = sock.recv(4)
            if not len_bytes or len(len_bytes) < 4:
                break
            length = int.from_bytes(len_bytes, byteorder="big")

            payload = bytearray()
            while len(payload) < length:
                chunk = sock.recv(length - len(payload))
                if not chunk:
                    break
                payload.extend(chunk)

            if len(payload) < length:
                break

            if msg_type == "KEYS":
                data = json.loads(payload.decode("utf-8"))
                event_type = data.get("type")
                key_name = data.get("key")
                if key_name and len(key_name) == 1:
                    key_name = key_name.lower()
                pg_key = KEY_MAP.get(key_name)
                if pg_key is not None:
                    if event_type == "keydown":
                        shared_keys[pg_key] = True
                        try:
                            pygame.event.post(
                                pygame.event.Event(
                                    pygame.KEYDOWN,
                                    key=pg_key,
                                    unicode=data.get("unicode", ""),
                                )
                            )
                        except Exception:
                            pass
                    elif event_type == "keyup":
                        shared_keys[pg_key] = False
                        try:
                            pygame.event.post(pygame.event.Event(pygame.KEYUP, key=pg_key))
                        except Exception:
                            pass

            elif msg_type == "MOUS":
                data = json.loads(payload.decode("utf-8"))
                shared_mouse_pos = (data.get("x", 0), data.get("y", 0))
                btn_left = data.get("btnLeft", False)
                btn_middle = data.get("btnMiddle", False)
                btn_right = data.get("btnRight", False)
                shared_mouse_buttons = (btn_left, btn_middle, btn_right)
                try:
                    event_type = data.get("type")
                    if event_type == "mousemove":
                        pygame.event.post(
                            pygame.event.Event(
                                pygame.MOUSEMOTION,
                                pos=shared_mouse_pos,
                                rel=(0, 0),
                                buttons=shared_mouse_buttons,
                            )
                        )
                    elif event_type == "mousedown":
                        button_idx = data.get("button", 0) + 1
                        pygame.event.post(
                            pygame.event.Event(
                                pygame.MOUSEBUTTONDOWN,
                                pos=shared_mouse_pos,
                                button=button_idx,
                            )
                        )
                    elif event_type == "mouseup":
                        button_idx = data.get("button", 0) + 1
                        pygame.event.post(
                            pygame.event.Event(
                                pygame.MOUSEBUTTONUP,
                                pos=shared_mouse_pos,
                                button=button_idx,
                            )
                        )
                except Exception:
                    pass

            elif msg_type == "CAMF":
                shared_webcam_frame = bytes(payload)

            elif msg_type == "CV2K":
                data = json.loads(payload.decode("utf-8"))
                key_code = data.get("key_code", -1)
                if key_code != -1:
                    cv2_keys.put(key_code)

            elif msg_type == "INPT":
                text = payload.decode("utf-8", errors="ignore")
                stdin_queue.put(text)

            elif msg_type == "QUIT":
                try:
                    pygame.event.post(pygame.event.Event(pygame.QUIT))
                except Exception:
                    pass

    except Exception as e:
        print(f"[Runner-IPC] Reader error: {e}", file=sys.__stderr__)
    finally:
        print("[Runner-IPC] Reader thread finished.", file=sys.__stderr__)
        try:
            sock.close()
        except Exception:
            pass


def writer_thread_func(sock):
    print("[Runner-IPC] Writer thread started.")
    try:
        while True:
            msg_type, payload = send_queue.get()
            if msg_type == b"EXIT":
                try:
                    sock.sendall(msg_type)
                    sock.sendall((0).to_bytes(4, byteorder="big"))
                except Exception:
                    pass
                break
            sock.sendall(msg_type)
            sock.sendall(len(payload).to_bytes(4, byteorder="big"))
            sock.sendall(payload)
    except Exception as e:
        print(f"[Runner-IPC] Writer error: {e}", file=sys.__stderr__)
    finally:
        print("[Runner-IPC] Writer thread finished.", file=sys.__stderr__)
        try:
            sock.close()
        except Exception:
            pass


sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
try:
    sock.connect(args.socket_path)
except Exception as e:
    print(f"Error connecting to server IPC socket: {e}")
    sys.exit(1)

reader_thread = threading.Thread(target=reader_thread_func, args=(sock,), daemon=True)
writer_thread = threading.Thread(target=writer_thread_func, args=(sock,), daemon=True)
reader_thread.start()
writer_thread.start()


# ----------------- Buffered stdout/stderr -----------------
class StreamToIPC:
    """Buffer log fragments and flush on newline or timer to reduce IPC chatter."""

    def __init__(self, original_stream, flush_interval=0.05):
        self.original_stream = original_stream
        self.flush_interval = flush_interval
        self._buf = []
        self._lock = threading.Lock()
        self._last_flush = time.time()
        self._stop = threading.Event()
        self._flusher = threading.Thread(target=self._flush_loop, daemon=True)
        self._flusher.start()

    def write(self, buf):
        if not buf:
            return 0
        self.original_stream.write(buf)
        self.original_stream.flush()
        with self._lock:
            self._buf.append(buf)
            if "\n" in buf or sum(len(x) for x in self._buf) > 2048:
                self._flush_unlocked()
        return len(buf)

    def _flush_unlocked(self):
        if not self._buf:
            return
        text = "".join(self._buf)
        self._buf.clear()
        self._last_flush = time.time()
        if text:
            send_queue.put((b"LOGS", text.encode("utf-8")))

    def _flush_loop(self):
        while not self._stop.wait(self.flush_interval):
            with self._lock:
                if self._buf and (time.time() - self._last_flush) >= self.flush_interval:
                    self._flush_unlocked()

    def flush(self):
        self.original_stream.flush()
        with self._lock:
            self._flush_unlocked()


sys.stdout = StreamToIPC(sys.stdout)
sys.stderr = StreamToIPC(sys.stderr)

# ----------------- Stdin / input() -----------------
import builtins

class CustomStdin:
    def readline(self, limit=-1):
        return stdin_queue.get()

    def read(self, size=-1):
        return stdin_queue.get()

    def fileno(self):
        raise OSError("Not a real file descriptor")

    def isatty(self):
        return True


sys.stdin = CustomStdin()

original_input = builtins.input


def custom_input(prompt=""):
    if prompt:
        sys.stdout.write(str(prompt))
        sys.stdout.flush()
    line = stdin_queue.get()
    if line.endswith("\r\n"):
        return line[:-2]
    if line.endswith("\n") or line.endswith("\r"):
        return line[:-1]
    return line


builtins.input = custom_input

# ----------------- Pygame patches -----------------
original_set_mode = pygame.display.set_mode


def custom_set_mode(size, flags=0, depth=0, display=0, vsync=0):
    surf = original_set_mode(size, flags, depth, display, vsync)
    print(f"[Pygame] 初始化視窗大小為: {size}")
    return surf


pygame.display.set_mode = custom_set_mode


def _encode_surface_jpeg(surf) -> bytes | None:
    """Encode pygame surface to JPEG via numpy + cv2 for controllable quality."""
    try:
        import cv2

        # pygame surface -> RGB array
        arr = pygame.surfarray.array3d(surf)
        # array3d is (w, h, 3); OpenCV wants (h, w, 3) BGR
        arr = np.transpose(arr, (1, 0, 2))
        bgr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
        # Optional downscale if very large
        h, w = bgr.shape[:2]
        max_w = 960
        if w > max_w:
            scale = max_w / float(w)
            bgr = cv2.resize(bgr, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)
        ok, jpeg = cv2.imencode(".jpg", bgr, [cv2.IMWRITE_JPEG_QUALITY, JPEG_QUALITY])
        if ok:
            return jpeg.tobytes()
    except Exception:
        # Fallback to pygame save
        try:
            buffer = io.BytesIO()
            pygame.image.save(surf, buffer, "JPEG")
            return buffer.getvalue()
        except Exception:
            return None
    return None


def send_pygame_frame():
    global last_pygame_frame_time
    now = time.time()
    qsize = send_queue.qsize()
    _adjust_frame_interval(qsize)
    if now - last_pygame_frame_time < _current_frame_interval:
        return
    if qsize > 3:
        return
    last_pygame_frame_time = now

    surf = pygame.display.get_surface()
    if not surf:
        return
    jpeg_bytes = _encode_surface_jpeg(surf)
    if jpeg_bytes:
        send_queue.put((b"FRME", jpeg_bytes))


original_flip = pygame.display.flip
original_update = pygame.display.update


def custom_flip():
    send_pygame_frame()
    original_flip()


def custom_update(*args, **kwargs):
    send_pygame_frame()
    original_update(*args, **kwargs)


pygame.display.flip = custom_flip
pygame.display.update = custom_update


def custom_key_get_pressed():
    return shared_keys


pygame.key.get_pressed = custom_key_get_pressed


def custom_mouse_get_pos():
    return shared_mouse_pos


pygame.mouse.get_pos = custom_mouse_get_pos


def custom_mouse_get_pressed(num_buttons=3):
    return shared_mouse_buttons


pygame.mouse.get_pressed = custom_mouse_get_pressed


def find_chinese_font_path():
    paths = [
        "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc",
        "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",
        "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc",
        "/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.ttc",
        "/usr/share/fonts/truetype/droid/DroidSansFallback.ttf",
        "/System/Library/Fonts/Supplemental/Songti.ttc",
        "/System/Library/Fonts/STHeiti Light.ttc",
        "/System/Library/Fonts/STHeiti Medium.ttc",
        "/System/Library/Fonts/PingFang.ttc",
        "/Library/Fonts/Arial Unicode.ttf",
        "C:\\Windows\\Fonts\\msjh.ttc",
        "C:\\Windows\\Fonts\\msjhbd.ttc",
        "C:\\Windows\\Fonts\\msyh.ttc",
        "C:\\Windows\\Fonts\\simsun.ttc",
    ]
    for p in paths:
        if os.path.exists(p):
            return p
    candidates = [
        "wenquanyizenhei", "wenquanyimicrohei", "wqy-zenhei", "wqy-microhei",
        "microsoftjhenghei", "pingfangtc", "songti", "stheitimedium",
        "arialunicodems", "droidsansfallback",
    ]
    for c in candidates:
        match = pygame.font.match_font(c)
        if match:
            return match
    return None


original_SysFont = pygame.font.SysFont


def custom_SysFont(name, size, bold=False, italic=False):
    name_lower = name.lower() if isinstance(name, str) else ""
    is_chinese_request = any(
        x in name_lower
        for x in ["jheng", "hei", "kai", "ming", "song", "cjk", "wqy", "unicode", "yahei", "wenquanyi"]
    ) or (name is None)
    if is_chinese_request:
        path = find_chinese_font_path()
        if path:
            try:
                return pygame.font.Font(path, size)
            except Exception:
                pass
    return original_SysFont(name, size, bold, italic)


pygame.font.SysFont = custom_SysFont

original_Font = pygame.font.Font


def custom_Font(filename, size):
    if filename is None:
        path = find_chinese_font_path()
        if path:
            try:
                return original_Font(path, size)
            except Exception:
                pass
    return original_Font(filename, size)


pygame.font.Font = custom_Font

# ----------------- OpenCV patches -----------------
import cv2  # noqa: E402


def custom_imshow(winname, mat):
    global last_cv2_frame_time
    now = time.time()
    qsize = send_queue.qsize()
    _adjust_frame_interval(qsize)
    if now - last_cv2_frame_time < _current_frame_interval:
        return
    if qsize > 3:
        return
    last_cv2_frame_time = now

    if mat is None:
        return
    try:
        frame = mat
        h, w = frame.shape[:2]
        max_w = 960
        if w > max_w:
            scale = max_w / float(w)
            frame = cv2.resize(frame, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)
        ret, jpeg = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, JPEG_QUALITY])
        if ret:
            send_queue.put((b"FRME", jpeg.tobytes()))
    except Exception:
        pass


cv2.imshow = custom_imshow


def custom_wait_key(delay=0):
    start_time = time.time()
    timeout = (delay / 1000.0) if delay > 0 else float("inf")
    while True:
        if not cv2_keys.empty():
            return cv2_keys.get()
        if time.time() - start_time >= timeout:
            break
        time.sleep(0.005)
    return -1


cv2.waitKey = custom_wait_key


class MockVideoCapture:
    def __init__(self, source=0):
        self.source = source
        self.is_opened = True
        self.frame_counter = 0

    def isOpened(self):
        return self.is_opened

    def read(self):
        global shared_webcam_frame
        if not self.is_opened:
            return False, None

        if shared_webcam_frame is not None:
            try:
                img_np = np.frombuffer(shared_webcam_frame, dtype=np.uint8)
                frame = cv2.imdecode(img_np, cv2.IMREAD_COLOR)
                if frame is not None:
                    return True, frame
            except Exception as e:
                print(f"[OpenCV-Mock] 解碼 Webcam 畫面錯誤: {e}", file=sys.__stderr__)

        width, height = 640, 480
        frame = np.zeros((height, width, 3), dtype=np.uint8)
        self.frame_counter += 1
        t = time.time()
        cx = int(320 + 120 * np.sin(t * 3))
        cy = int(240 + 80 * np.cos(t * 2))
        cv2.circle(frame, (cx, cy), 40, (150, 80, 230), -1)
        cv2.putText(
            frame,
            "Webcam Connecting...",
            (160, 240),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.8,
            (240, 240, 240),
            2,
        )
        cv2.putText(
            frame,
            "Please enable webcam in browser",
            (130, 280),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            (200, 200, 200),
            1,
        )
        return True, frame

    def release(self):
        self.is_opened = False

    def get(self, propId):
        if propId == cv2.CAP_PROP_FRAME_WIDTH:
            return 640
        if propId == cv2.CAP_PROP_FRAME_HEIGHT:
            return 480
        if propId == cv2.CAP_PROP_FPS:
            return 30
        return 0

    def set(self, propId, value):
        return True


original_VideoCapture = cv2.VideoCapture


def custom_VideoCapture(source=0):
    if isinstance(source, int) or source == "0" or source == 0:
        print("[OpenCV] 攔截實體鏡頭 VideoCapture(0) 啟用瀏覽器 Webcam 串流。")
        return MockVideoCapture(source)
    print(f"[OpenCV] 載入影片檔案: {source}")
    return original_VideoCapture(source)


cv2.VideoCapture = custom_VideoCapture


# ----------------- Shared models + project setup -----------------
def ensure_shared_models(project_dir: str, models_dir: str | None) -> None:
    """Symlink common MediaPipe models from shared models/ if missing in project."""
    if not models_dir or not os.path.isdir(models_dir):
        return
    known = (
        "hand_landmarker.task",
        "face_landmarker.task",
        "pose_landmarker.task",
        "gesture_recognizer.task",
    )
    for name in known:
        dest = os.path.join(project_dir, name)
        src = os.path.join(models_dir, name)
        if os.path.exists(dest) or not os.path.isfile(src):
            continue
        try:
            os.symlink(src, dest)
            print(f"[Runner] 連結共用模型: {name} -> {src}")
        except OSError:
            # Fallback: copy is too heavy; tell user models_dir path via env
            try:
                # Hard-link if same filesystem
                os.link(src, dest)
                print(f"[Runner] 硬連結共用模型: {name}")
            except OSError as e:
                print(f"[Runner] 無法連結模型 {name}: {e}（請將檔案放在專案目錄或 models/）")


import runpy

project_path = os.path.abspath(args.project_path)
project_dir = os.path.dirname(project_path)

print(f"[Runner] 切換工作目錄至: {project_dir}")
os.chdir(project_dir)
sys.path.insert(0, project_dir)

models_dir = args.models_dir
if not models_dir:
    # Default: <repo>/models next to runner.py
    models_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
ensure_shared_models(project_dir, models_dir)

# Keras dummy assets → temp session dir (do not pollute student project)
try:
    script_content = open(project_path, "r", errors="ignore").read()
    if "keras_model.h5" in script_content:
        session_tmp = tempfile.mkdtemp(prefix="telearn_keras_")
        h5_path = os.path.join(project_dir, "keras_model.h5")
        labels_path = os.path.join(project_dir, "labels.txt")
        tmp_h5 = os.path.join(session_tmp, "keras_model.h5")
        tmp_labels = os.path.join(session_tmp, "labels.txt")

        need_h5 = not os.path.exists(h5_path)
        need_labels = not os.path.exists(labels_path)

        if need_h5 or need_labels:
            print(f"[Runner] Keras demo 資源寫入暫存目錄: {session_tmp}")

        if need_h5:
            print("[Runner] 偵測到 Keras 專案且缺乏 keras_model.h5，嘗試自動生成 Dummy 模型...")
            try:
                try:
                    import keras
                except ImportError:
                    from tensorflow import keras  # type: ignore

                model = keras.models.Sequential(
                    [
                        keras.layers.Input(shape=(224, 224, 3)),
                        keras.layers.Flatten(),
                        keras.layers.Dense(3, activation="softmax"),
                    ]
                )
                model.save(tmp_h5)
                try:
                    os.symlink(tmp_h5, h5_path)
                except OSError:
                    # Last resort: only if project dir is writable and no existing file
                    if not os.path.exists(h5_path):
                        import shutil

                        shutil.copy2(tmp_h5, h5_path)
                print("[Runner] 成功生成 dummy keras_model.h5")
            except Exception as ex:
                print(f"[Runner] 無法自動生成模型（可能未安裝 tensorflow/keras）：{ex}")

        if need_labels:
            print("[Runner] 偵測到缺乏 labels.txt，自動生成預設分類標籤...")
            try:
                with open(tmp_labels, "w", encoding="utf-8") as f:
                    f.write("0 背景\n1 喝水中\n2 其他動作\n")
                try:
                    os.symlink(tmp_labels, labels_path)
                except OSError:
                    if not os.path.exists(labels_path):
                        with open(labels_path, "w", encoding="utf-8") as f:
                            f.write("0 背景\n1 喝水中\n2 其他動作\n")
                print("[Runner] 成功生成 labels.txt")
            except Exception as ex:
                print(f"[Runner] 無法自動生成 labels.txt：{ex}")
except Exception as e:
    print(f"[Runner] 檢查/生成 Keras 依賴資源時出錯: {e}")

try:
    print(f"[Runner] 開始執行學生專案: {os.path.basename(project_path)}")
    runpy.run_path(project_path, run_name="__main__")
except SystemExit as e:
    print(f"[Runner] 專案正常結束 (SystemExit: {e})")
except Exception:
    import traceback

    print("[Runner] 專案執行出錯:")
    traceback.print_exc(file=sys.stderr)
finally:
    print("[Runner] 清理資源，結束執行。")
    try:
        sys.stdout.flush()
        sys.stderr.flush()
    except Exception:
        pass
    send_queue.put((b"EXIT", b""))
    time.sleep(0.5)