ui.py

"""Small native UI for the face-tracking application."""

from __future__ import annotations

from io import BytesIO
import tkinter as tk
from pathlib import Path
from tkinter import filedialog, messagebox, ttk

from PIL import Image, ImageDraw, ImageTk, UnidentifiedImageError


# Reserve vertical room for music controls in the same desktop window.
PREVIEW_MAX_SIZE = (300, 300)


def render_preview_image(jpeg: bytes, faces: tuple[object, ...]) -> Image.Image:
    """Decode one SDK JPEG and overlay the associated face boxes."""

    with Image.open(BytesIO(jpeg)) as source:
        image = source.convert("RGB")

    drawing = ImageDraw.Draw(image)
    for index, face in enumerate(faces, start=1):
        x = int(getattr(face, "x"))
        y = int(getattr(face, "y"))
        width = int(getattr(face, "width"))
        height = int(getattr(face, "height"))
        score = int(getattr(face, "score"))
        left = x - width // 2
        top = y - height // 2
        right = left + width
        bottom = top + height
        drawing.rectangle((left, top, right, bottom), outline="#ff3b30", width=3)
        drawing.text((left + 4, max(0, top - 18)), f"face {index}: {score}", fill="#ff3b30")

    image.thumbnail(PREVIEW_MAX_SIZE, Image.Resampling.LANCZOS)
    return image


class PreviewWindow:
    """Keep pairing, preview start/stop, and application status in one local window."""

    def __init__(self) -> None:
        self._is_open = True
        self._pairing_code: str | None = None
        self._preview_action: str | None = None
        self._music_action: str | None = None
        self._music_path: Path | None = None
        self._online = False
        self._preview_running = False
        self._music_running = False
        self._root = tk.Tk()
        self._root.title("人脸跟随应用 - WatcheRobot")
        self._root.geometry("600x1000")
        self._root.minsize(520, 800)
        self._root.protocol("WM_DELETE_WINDOW", self.close)

        container = ttk.Frame(self._root, padding=32)
        container.pack(fill="both", expand=True)
        ttk.Label(
            container,
            text="人脸跟随应用",
            font=("Microsoft YaHei UI", 22, "bold"),
        ).pack(anchor="w")
        ttk.Label(
            container,
            text="此项目只读取预览帧与人脸框遥测，不进行身份识别。",
            wraplength=500,
        ).pack(anchor="w", pady=(10, 0))

        pairing = ttk.LabelFrame(container, text="连接机器人", padding=16)
        pairing.pack(fill="x", pady=(22, 0))
        ttk.Label(pairing, text="输入机器人屏幕上的 6 位配对码。").pack(anchor="w")
        controls = ttk.Frame(pairing)
        controls.pack(fill="x", pady=(10, 0))
        self._pairing_code_var = tk.StringVar()
        self._pairing_entry = ttk.Entry(controls, textvariable=self._pairing_code_var, width=18)
        self._pairing_entry.pack(side="left")
        self._connect_button = ttk.Button(controls, text="连接", command=self._request_pairing)
        self._connect_button.pack(side="left", padx=(10, 0))
        self._connection_var = tk.StringVar(value="尚未连接")
        ttk.Label(pairing, textvariable=self._connection_var).pack(anchor="w", pady=(12, 0))

        preview = ttk.LabelFrame(container, text="相机预览", padding=16)
        preview.pack(fill="x", pady=(18, 0))
        preview_controls = ttk.Frame(preview)
        preview_controls.pack(anchor="w")
        self._start_button = ttk.Button(preview_controls, text="开始预览", command=self._request_start)
        self._start_button.pack(side="left")
        self._stop_button = ttk.Button(preview_controls, text="停止预览", command=self._request_stop)
        self._stop_button.pack(side="left", padx=(10, 0))
        self._preview_var = tk.StringVar(value="等待开始预览")
        ttk.Label(preview, textvariable=self._preview_var, wraplength=500).pack(anchor="w", pady=(12, 0))
        self._preview_image = ttk.Label(
            preview,
            text="开始预览后，相机画面会显示在这里。",
            anchor="center",
            justify="center",
            width=58,
        )
        self._preview_image.pack(fill="both", expand=True, pady=(12, 0))
        self._preview_photo: ImageTk.PhotoImage | None = None

        music = ttk.LabelFrame(container, text="播放音乐", padding=16)
        music.pack(fill="x", pady=(18, 0))
        ttk.Label(
            music,
            text="播放时会自动停止相机预览，并让机器人显示音乐表情。",
            wraplength=500,
        ).pack(anchor="w")
        self._music_path_var = tk.StringVar(value="尚未选择音频文件")
        ttk.Label(music, textvariable=self._music_path_var, wraplength=500).pack(
            anchor="w", pady=(10, 0)
        )
        music_controls = ttk.Frame(music)
        music_controls.pack(anchor="w", pady=(10, 0))
        self._select_music_button = ttk.Button(
            music_controls,
            text="选择音频",
            command=self._request_select_music,
        )
        self._select_music_button.pack(side="left")
        self._play_music_button = ttk.Button(
            music_controls,
            text="播放音乐",
            command=self._request_play_music,
        )
        self._play_music_button.pack(side="left", padx=(10, 0))
        self._stop_music_button = ttk.Button(
            music_controls,
            text="停止音乐",
            command=self._request_stop_music,
        )
        self._stop_music_button.pack(side="left", padx=(10, 0))
        self._music_var = tk.StringVar(value="选择音频后即可播放到机器人扬声器。")
        ttk.Label(music, textvariable=self._music_var, wraplength=500).pack(
            anchor="w", pady=(12, 0)
        )

        ttk.Button(container, text="关闭", command=self.close).pack(anchor="e", pady=(20, 0))
        self._set_preview_buttons(online=False, running=False)
        self._set_music_buttons(online=False, running=False)

    @property
    def is_open(self) -> bool:
        return self._is_open

    def pump(self) -> None:
        if not self._is_open:
            return
        try:
            self._root.update_idletasks()
            self._root.update()
        except tk.TclError:
            self._is_open = False

    def take_pairing_code(self) -> str | None:
        pairing_code = self._pairing_code
        self._pairing_code = None
        return pairing_code

    def take_preview_action(self) -> str | None:
        action = self._preview_action
        self._preview_action = None
        return action

    def take_music_action(self) -> tuple[str, Path | None] | None:
        action = self._music_action
        self._music_action = None
        if action is None:
            return None
        return action, self._music_path

    def update_connection(self, status: object) -> None:
        if not isinstance(status, dict):
            self.show_status_unavailable()
            return
        online = status.get("online") is True
        self._online = online
        state = str(status.get("state") or "")
        if online:
            name = str(status.get("name") or "机器人")
            self._connection_var.set(f"已连接：{name}")
            self._set_pairing_enabled(False)
            self._set_preview_buttons(online=True, running=self._preview_running)
            self._set_music_buttons(online=True, running=self._music_running)
            return
        if state in {"discovering", "connecting", "reconnecting"}:
            self._connection_var.set("正在连接，请稍候…")
            self._set_pairing_enabled(False)
        else:
            self._connection_var.set("尚未连接，请输入设备上的 6 位配对码。")
            self._set_pairing_enabled(True)
        self._set_preview_buttons(online=False, running=False)
        self._set_music_buttons(online=False, running=False)

    def show_pairing_error(self) -> None:
        self._connection_var.set("连接未成功，请确认配对码和 Wi-Fi 后重试。")
        self._set_pairing_enabled(True)

    def show_status_unavailable(self) -> None:
        self._online = False
        self._connection_var.set("暂时无法检查连接状态，请稍后重试。")
        self._set_pairing_enabled(True)
        self._set_preview_buttons(online=False, running=False)
        self._set_music_buttons(online=False, running=False)

    def show_preview_running(self) -> None:
        self._preview_running = True
        self._preview_var.set("预览已启动，正在等待设备帧…")
        self._set_preview_buttons(online=True, running=True)

    def show_preview_waiting(self) -> None:
        self._preview_var.set("预览已启动，但暂未收到完整帧。")

    def show_preview_frame(
        self,
        *,
        jpeg: bytes,
        faces: tuple[object, ...],
        frame_count: int,
        face_count: int,
        age_ms: int,
        inference_ms: int | None,
        dropped_frames: int,
    ) -> None:
        inference = "未知" if inference_ms is None else f"{inference_ms} ms"
        self._preview_var.set(
            f"帧 {frame_count}：检测到 {face_count} 张脸；帧龄 {age_ms} ms；"
            f"推理 {inference}；丢帧 {dropped_frames}。"
        )
        try:
            self._preview_photo = self._render_preview(jpeg, faces)
        except (OSError, UnidentifiedImageError) as error:
            self._preview_image.configure(text=f"无法显示当前相机帧：{error}", image="")
            return
        self._preview_image.configure(image=self._preview_photo, text="")

    @staticmethod
    def _render_preview(
        jpeg: bytes,
        faces: tuple[object, ...],
    ) -> ImageTk.PhotoImage:
        return ImageTk.PhotoImage(render_preview_image(jpeg, faces))

    def show_preview_stopping(self) -> None:
        self._preview_var.set("正在停止预览…")
        self._set_preview_buttons(online=True, running=True)

    def show_preview_stopped(self) -> None:
        self._preview_running = False
        self._preview_var.set("预览已停止；最后一帧会保存到 artifacts/latest.jpg。")
        self._set_preview_buttons(online=True, running=False)

    def show_preview_error(self, message: str) -> None:
        self._preview_running = False
        self._preview_var.set(f"预览不可用：{message}")
        self._set_preview_buttons(online=True, running=False)

    def show_music_preparing(self, source: Path) -> None:
        self._music_running = True
        self._music_var.set(f"正在准备《{source.name}》…")
        self._set_preview_buttons(online=True, running=False)
        self._set_music_buttons(online=True, running=True)

    def show_music_playing(self, source: Path, duration_seconds: float) -> None:
        self._music_var.set(
            f"正在播放《{source.name}》；机器人保持音乐表情，预计 {duration_seconds:.1f} 秒。"
        )

    def show_music_stopping(self) -> None:
        self._music_var.set("正在停止音乐…")

    def show_music_stopped(self, message: str = "音乐已停止，机器人已回到待机表情。") -> None:
        self._music_running = False
        self._music_var.set(message)
        self._set_preview_buttons(online=self._online, running=False)
        self._set_music_buttons(online=self._online, running=False)

    def show_music_error(self, message: str) -> None:
        self._music_running = False
        self._music_var.set(f"音乐不可用：{message}")
        self._set_preview_buttons(online=self._online, running=False)
        self._set_music_buttons(online=self._online, running=False)

    def _request_pairing(self) -> None:
        pairing_code = self._pairing_code_var.get().strip()
        if len(pairing_code) != 6 or not pairing_code.isdigit():
            self._connection_var.set("请输入机器人显示的 6 位数字配对码。")
            return
        self._pairing_code_var.set("")
        self._pairing_code = pairing_code
        self._connection_var.set("正在连接，请稍候…")
        self._set_pairing_enabled(False)

    def _request_start(self) -> None:
        self._preview_action = "start"

    def _request_stop(self) -> None:
        self._preview_action = "stop"

    def _request_select_music(self) -> None:
        selected = filedialog.askopenfilename(
            parent=self._root,
            title="选择要播放到机器人的音频",
            filetypes=[
                ("音频文件", "*.wav *.mp3 *.m4a *.aac *.flac *.ogg"),
                ("所有文件", "*.*"),
            ],
        )
        if not selected:
            return
        self._music_path = Path(selected)
        self._music_path_var.set(f"已选择：{self._music_path.name}")
        self._music_var.set("点击“播放音乐”后，相机预览会停止。")
        self._set_music_buttons(online=self._online, running=False)

    def _request_play_music(self) -> None:
        if self._music_path is None:
            self._music_var.set("请先选择音频文件。")
            return
        self._music_action = "play"

    def _request_stop_music(self) -> None:
        if not messagebox.askyesno(
            "停止音乐",
            "确定停止当前音乐吗？",
            parent=self._root,
        ):
            return
        self._music_action = "stop"

    def _set_pairing_enabled(self, enabled: bool) -> None:
        if not self._is_open:
            return
        state = "normal" if enabled else "disabled"
        try:
            self._pairing_entry.configure(state=state)
            self._connect_button.configure(state=state)
        except tk.TclError:
            self._is_open = False

    def _set_preview_buttons(self, *, online: bool, running: bool) -> None:
        if not self._is_open:
            return
        try:
            music_running = getattr(self, "_music_running", False)
            self._start_button.configure(state="disabled" if running or music_running or not online else "normal")
            self._stop_button.configure(state="normal" if running and not music_running else "disabled")
        except tk.TclError:
            self._is_open = False

    def _set_music_buttons(self, *, online: bool, running: bool) -> None:
        if not self._is_open:
            return
        try:
            has_music = getattr(self, "_music_path", None) is not None
            self._select_music_button.configure(state="disabled" if running or not online else "normal")
            self._play_music_button.configure(
                state="disabled" if running or not online or not has_music else "normal"
            )
            self._stop_music_button.configure(state="normal" if running else "disabled")
        except tk.TclError:
            self._is_open = False

    def close(self) -> None:
        if not self._is_open:
            return
        self._is_open = False
        try:
            self._root.destroy()
        except tk.TclError:
            pass