app.py

#!/usr/bin/env python3
"""macOS GUI for Wi‑Fi Roomba control — smart commands, timeouts, auto-reconnect."""

from __future__ import annotations

import atexit
import json
import threading
import tkinter as tk
from pathlib import Path
from tkinter import messagebox, ttk

from automation import AutomationEngine, AutomationSettings
from mission_panel import MissionPanel
from roomba_driver import (
    DiscoveredRobot,
    RoombaDriver,
    discover_at_ip,
    discover_robots,
    fetch_password,
)

CONFIG_PATH = Path(__file__).resolve().parent / "config.json"
BUSY_HARD_TIMEOUT_MS = 25_000


class RoombaApp(tk.Tk):
    def __init__(self) -> None:
        super().__init__()
        self.title("Roomba Wi‑Fi Controller")
        self.geometry("1000x780")
        self.minsize(860, 700)

        self.driver = RoombaDriver()
        self.driver.register_on_disconnect(self._on_driver_disconnect)
        self.automation = AutomationEngine()
        self._discovered: list[DiscoveredRobot] = []
        self._closing = False
        self._busy = False
        self._busy_job: str | None = None
        self._key_armed = True  # ignore OS key-repeat floods
        self._auto_tick = 0
        self._auto_running = False

        self._build_ui()
        self._bind_keys()
        self._load_config()

        self.protocol("WM_DELETE_WINDOW", self.on_quit)
        atexit.register(self._cleanup)
        self.after(200, self._tick_status)
        self.after(400, self._auto_connect_if_possible)

    def _build_ui(self) -> None:
        pad = {"padx": 10, "pady": 6}

        body = ttk.Panedwindow(self, orient=tk.HORIZONTAL)
        body.pack(fill="both", expand=True, padx=8, pady=8)
        left = ttk.Frame(body)
        right = ttk.Frame(body)
        body.add(left, weight=1)
        body.add(right, weight=2)

        conn = ttk.LabelFrame(left, text="Connection")
        conn.pack(fill="x", **pad)

        ttk.Label(conn, text="Robot").grid(row=0, column=0, sticky="w", padx=8, pady=6)
        self.robot_var = tk.StringVar()
        self.robot_combo = ttk.Combobox(
            conn, textvariable=self.robot_var, width=40, state="readonly"
        )
        self.robot_combo.grid(row=0, column=1, columnspan=2, sticky="ew", padx=8, pady=6)
        self.robot_combo.bind("<<ComboboxSelected>>", self._on_robot_selected)

        self.discover_btn = ttk.Button(conn, text="Discover", command=self.discover)
        self.discover_btn.grid(row=0, column=3, padx=8, pady=6)

        ttk.Label(conn, text="IP").grid(row=1, column=0, sticky="w", padx=8, pady=4)
        self.ip_var = tk.StringVar()
        ttk.Entry(conn, textvariable=self.ip_var, width=22).grid(
            row=1, column=1, sticky="w", padx=8, pady=4
        )

        ttk.Label(conn, text="BLID").grid(row=2, column=0, sticky="w", padx=8, pady=4)
        self.blid_var = tk.StringVar()
        ttk.Entry(conn, textvariable=self.blid_var, width=40).grid(
            row=2, column=1, columnspan=2, sticky="ew", padx=8, pady=4
        )

        ttk.Label(conn, text="Password").grid(row=3, column=0, sticky="w", padx=8, pady=4)
        self.password_var = tk.StringVar()
        ttk.Entry(conn, textvariable=self.password_var, width=40, show="•").grid(
            row=3, column=1, columnspan=2, sticky="ew", padx=8, pady=4
        )
        self.pw_btn = ttk.Button(conn, text="Get password", command=self.get_password)
        self.pw_btn.grid(row=3, column=3, padx=8, pady=4)

        self.connect_btn = ttk.Button(conn, text="Connect", command=self.toggle_connect)
        self.connect_btn.grid(row=4, column=1, sticky="w", padx=8, pady=8)
        self.auto_var = tk.BooleanVar(value=True)
        ttk.Checkbutton(
            conn, text="Auto-connect on launch", variable=self.auto_var
        ).grid(row=4, column=2, sticky="w", padx=8, pady=8)
        conn.columnconfigure(1, weight=1)

        ttk.Label(
            conn,
            text=(
                "First time: Discover → dock robot → hold HOME until Wi‑Fi flashes "
                "→ Get password → Connect"
            ),
            wraplength=500,
            font=("", 10),
        ).grid(row=5, column=0, columnspan=4, sticky="w", padx=8, pady=(0, 8))

        actions = ttk.LabelFrame(left, text="Actions")
        actions.pack(fill="x", **pad)

        btn_row = ttk.Frame(actions)
        btn_row.pack(fill="x", padx=8, pady=8)
        self.btn_clean = ttk.Button(
            btn_row, text="Clean (W)", command=lambda: self._run_smart(self.driver.smart_clean)
        )
        self.btn_pause = ttk.Button(
            btn_row, text="Pause (P)", command=lambda: self._run_smart(self.driver.smart_pause)
        )
        self.btn_dock = ttk.Button(
            btn_row, text="Go home (Space)", command=lambda: self._run_smart(self.driver.smart_dock)
        )
        self.btn_stop = ttk.Button(
            btn_row, text="Stop (S)", command=lambda: self._run_smart(self.driver.smart_stop)
        )
        for i, b in enumerate(
            (self.btn_clean, self.btn_pause, self.btn_dock, self.btn_stop)
        ):
            b.grid(row=0, column=i, sticky="ew", padx=4)
            btn_row.columnconfigure(i, weight=1)

        ttk.Label(
            actions,
            text="Commands adapt to state (e.g. W resumes if paused; Space pauses then docks).",
            font=("", 10),
        ).pack(anchor="w", padx=8, pady=(0, 8))

        auto = ttk.LabelFrame(left, text="Automations (app must stay open + connected)")
        auto.pack(fill="x", **pad)

        # Schedule row
        self.sched_enabled = tk.BooleanVar(value=False)
        ttk.Checkbutton(
            auto, text="Clean every day at", variable=self.sched_enabled, command=self._on_auto_changed
        ).grid(row=0, column=0, sticky="w", padx=8, pady=4)
        self.sched_time = tk.StringVar(value="09:00")
        ttk.Entry(auto, textvariable=self.sched_time, width=7).grid(
            row=0, column=1, sticky="w", padx=4, pady=4
        )
        ttk.Label(auto, text="(HH:MM)", font=("", 10)).grid(
            row=0, column=2, sticky="w", padx=4, pady=4
        )

        # Battery start
        self.bat_start_enabled = tk.BooleanVar(value=False)
        ttk.Checkbutton(
            auto,
            text="Start clean when battery reaches",
            variable=self.bat_start_enabled,
            command=self._on_auto_changed,
        ).grid(row=1, column=0, sticky="w", padx=8, pady=4)
        self.bat_start_pct = tk.StringVar(value="50")
        ttk.Entry(auto, textvariable=self.bat_start_pct, width=5).grid(
            row=1, column=1, sticky="w", padx=4, pady=4
        )
        ttk.Label(auto, text="% (while docked)", font=("", 10)).grid(
            row=1, column=2, sticky="w", padx=4, pady=4
        )

        # Battery dock
        self.bat_dock_enabled = tk.BooleanVar(value=False)
        ttk.Checkbutton(
            auto,
            text="Go home when battery falls to",
            variable=self.bat_dock_enabled,
            command=self._on_auto_changed,
        ).grid(row=2, column=0, sticky="w", padx=8, pady=4)
        self.bat_dock_pct = tk.StringVar(value="15")
        ttk.Entry(auto, textvariable=self.bat_dock_pct, width=5).grid(
            row=2, column=1, sticky="w", padx=4, pady=4
        )
        ttk.Label(auto, text="% (while cleaning)", font=("", 10)).grid(
            row=2, column=2, sticky="w", padx=4, pady=4
        )

        self.auto_status = ttk.Label(auto, text="Automations off", font=("", 10))
        self.auto_status.grid(row=3, column=0, columnspan=3, sticky="w", padx=8, pady=(2, 8))

        status = ttk.LabelFrame(left, text="Status")
        status.pack(fill="both", expand=True, **pad)

        self.conn_label = ttk.Label(status, text="Connection: Disconnected", font=("", 12, "bold"))
        self.conn_label.pack(anchor="w", padx=8, pady=4)
        self.state_label = ttk.Label(status, text="State: —", font=("", 14, "bold"))
        self.state_label.pack(anchor="w", padx=8, pady=6)
        self.name_label = ttk.Label(status, text="Name: —")
        self.name_label.pack(anchor="w", padx=8, pady=2)
        self.battery_label = ttk.Label(status, text="Battery: —")
        self.battery_label.pack(anchor="w", padx=8, pady=2)
        self.detail_label = ttk.Label(status, text="Detail: —")
        self.detail_label.pack(anchor="w", padx=8, pady=2)
        self.pose_label = ttk.Label(status, text="Pose: —")
        self.pose_label.pack(anchor="w", padx=8, pady=2)
        self.cmd_label = ttk.Label(status, text="Last action: —")
        self.cmd_label.pack(anchor="w", padx=8, pady=2)
        self.error_label = ttk.Label(status, text="", foreground="#b00020", wraplength=360)
        self.error_label.pack(anchor="w", padx=8, pady=6)

        map_frame = ttk.LabelFrame(right, text="Mission view")
        map_frame.pack(fill="both", expand=True, **pad)
        self.mission_panel = MissionPanel(map_frame, self.driver.map_tracker)
        self.mission_panel.pack(fill="both", expand=True, padx=4, pady=4)

    def _bind_keys(self) -> None:
        mapping = {
            "w": self.driver.smart_clean,
            "s": self.driver.smart_stop,
            "r": self.driver.smart_resume,
            "p": self.driver.smart_pause,
            "space": self.driver.smart_dock,
        }

        def press(event: tk.Event) -> str:
            if not self._key_armed:
                return "break"
            key = event.keysym.lower()
            fn = mapping.get(key)
            if not fn:
                return "break"
            self._key_armed = False
            self._run_smart(fn)
            return "break"

        def release(_event: tk.Event) -> str:
            self._key_armed = True
            return "break"

        for key in ("w", "a", "s", "d", "r", "p", "W", "S", "R", "P"):
            self.bind(f"<KeyPress-{key}>", press)
            self.bind(f"<KeyRelease-{key}>", release)
        self.bind("<KeyPress-space>", press)
        self.bind("<KeyRelease-space>", release)
        self.bind("<Escape>", lambda _e: self.on_quit())

    def _set_busy(self, job: str, message: str) -> bool:
        if self._busy:
            self.error_label.configure(text=f"Busy: {self._busy_job}")
            return False
        self._busy = True
        self._busy_job = job
        self.error_label.configure(text=message)
        self._set_conn_controls(False)
        self.after(BUSY_HARD_TIMEOUT_MS, lambda: self._force_clear_busy(job))
        return True

    def _clear_busy(self) -> None:
        self._busy = False
        self._busy_job = None
        self._set_conn_controls(True)

    def _force_clear_busy(self, job: str) -> None:
        if self._busy and self._busy_job == job:
            self._clear_busy()
            self.error_label.configure(text=f"{job} took too long — unlocked UI")

    def _set_conn_controls(self, enabled: bool) -> None:
        state = "normal" if enabled else "disabled"
        for w in (self.discover_btn, self.pw_btn, self.connect_btn):
            w.configure(state=state)

    def _set_action_controls(self, enabled: bool) -> None:
        state = "normal" if enabled else "disabled"
        for w in (self.btn_clean, self.btn_pause, self.btn_dock, self.btn_stop):
            w.configure(state=state)

    def discover(self) -> None:
        if not self._set_busy("Discover", "Discovering on LAN…"):
            return

        def work() -> None:
            err = ""
            robots: list[DiscoveredRobot] = []
            try:
                robots = discover_robots()
                ip = self.ip_var.get().strip()
                if not robots and ip:
                    one = discover_at_ip(ip)
                    if one:
                        robots = [one]
            except Exception as exc:
                err = str(exc)

            def done() -> None:
                self._clear_busy()
                if err:
                    self.error_label.configure(text=err)
                    return
                self._discovered = robots
                labels = [f"{r.name}  ({r.ip})" for r in robots]
                self.robot_combo["values"] = labels
                if robots:
                    self.robot_combo.current(0)
                    self._apply_robot(robots[0])
                    self.error_label.configure(text=f"Found {len(robots)} robot(s)")
                    self._save_config()
                else:
                    self.error_label.configure(
                        text="No robots found — check Wi‑Fi, or type IP from your router"
                    )

            self.after(0, done)

        threading.Thread(target=work, daemon=True).start()

    def _on_robot_selected(self, _event: tk.Event | None = None) -> None:
        idx = self.robot_combo.current()
        if 0 <= idx < len(self._discovered):
            self._apply_robot(self._discovered[idx])

    def _apply_robot(self, robot: DiscoveredRobot) -> None:
        self.ip_var.set(robot.ip)
        self.blid_var.set(robot.blid)
        if robot.password:
            self.password_var.set(robot.password)

    def get_password(self) -> None:
        ip = self.ip_var.get().strip()
        if not ip:
            messagebox.showwarning("IP required", "Enter the Roomba IP first (or Discover).")
            return
        if not self._set_busy("Get password", "Requesting password… hold HOME / flashing Wi‑Fi"):
            return

        def work() -> None:
            pwd = None
            err = ""
            blid = ""
            try:
                bot = discover_at_ip(ip)
                if bot:
                    blid = bot.blid
                pwd = fetch_password(ip)
            except Exception as exc:
                err = str(exc)

            def done() -> None:
                self._clear_busy()
                if blid:
                    self.blid_var.set(blid)
                if err:
                    self.error_label.configure(text=err)
                    return
                if not pwd:
                    self.error_label.configure(
                        text="No password — hold HOME until Wi‑Fi LED flashes, then retry quickly"
                    )
                    return
                self.password_var.set(pwd)
                self.error_label.configure(text="Password saved")
                self._save_config()
                self.toggle_connect()

            self.after(0, done)

        threading.Thread(target=work, daemon=True).start()

    def toggle_connect(self) -> None:
        if self.driver.connected or self.driver.status.reconnecting:
            self._disconnect()
            return

        ip = self.ip_var.get().strip()
        blid = self.blid_var.get().strip()
        password = self.password_var.get().strip()
        if not ip or not blid or not password:
            messagebox.showwarning(
                "Missing credentials",
                "Need IP, BLID, and password.\nUse Discover + Get password first.",
            )
            return
        if not self._set_busy("Connect", "Connecting…"):
            return

        def work() -> None:
            err = ""
            try:
                self.driver.connect(ip, blid, password)
            except Exception as exc:
                err = str(exc)

            def done() -> None:
                self._clear_busy()
                if err:
                    self.error_label.configure(text=err)
                    self.connect_btn.configure(text="Connect")
                    return
                self.connect_btn.configure(text="Disconnect")
                self._save_config()
                self.error_label.configure(text="")

            self.after(0, done)

        threading.Thread(target=work, daemon=True).start()

    def _auto_connect_if_possible(self) -> None:
        if not self.auto_var.get():
            return
        if self.driver.connected or self._busy:
            return
        ip = self.ip_var.get().strip()
        blid = self.blid_var.get().strip()
        password = self.password_var.get().strip()
        if ip and blid and password:
            self.toggle_connect()

    def _disconnect(self) -> None:
        try:
            self.driver.disconnect()
        except Exception:
            pass
        self.connect_btn.configure(text="Connect")
        self.error_label.configure(text="")

    def _run_smart(self, fn) -> None:
        if not self.driver.connected:
            self.error_label.configure(text="Not connected")
            return

        def work() -> None:
            msg = ""
            err = ""
            try:
                msg = fn()
            except Exception as exc:
                err = str(exc)

            def done() -> None:
                if err:
                    self.error_label.configure(text=err)
                elif msg:
                    self.cmd_label.configure(text=f"Last action: {msg}")
                    self.error_label.configure(text="")

            self.after(0, done)

        threading.Thread(target=work, daemon=True).start()

    def _on_driver_disconnect(self) -> None:
        def ui() -> None:
            if not self.driver.status.reconnecting:
                self.connect_btn.configure(text="Connect")

        self.after(0, ui)

    def _tick_status(self) -> None:
        if self._closing:
            return
        st = self.driver.status

        if st.connecting:
            self.conn_label.configure(text=f"Connection: Connecting ({st.ip})…")
        elif st.reconnecting:
            self.conn_label.configure(text=f"Connection: Reconnecting ({st.ip})…")
            self.connect_btn.configure(text="Disconnect")
        elif st.connected:
            self.conn_label.configure(text=f"Connection: Connected ({st.ip})")
            self.connect_btn.configure(text="Disconnect")
        else:
            self.conn_label.configure(text="Connection: Disconnected")
            if not self._busy:
                self.connect_btn.configure(text="Connect")

        self.state_label.configure(text=f"State: {st.state_label}")
        self.name_label.configure(text=f"Name: {st.name or '—'}")
        if st.battery_pct is None:
            self.battery_label.configure(text="Battery: —")
        else:
            self.battery_label.configure(text=f"Battery: {st.battery_pct}%")
        self.detail_label.configure(text=f"Detail: phase={st.phase}  cycle={st.cycle}")
        self.pose_label.configure(
            text=f"Pose: {st.pose_xy}" + (f"  ·  {st.sku}" if st.sku else "")
        )
        self.cmd_label.configure(text=f"Last action: {st.last_command}")
        if st.error:
            self.error_label.configure(text=st.error)

        self._set_action_controls(st.connected and not st.connecting)

        try:
            self.mission_panel.update(
                connected=st.connected,
                pose_supported=st.pose_supported,
                battery_pct=st.battery_pct,
                phase=st.phase,
                cycle=st.cycle,
                state_label=st.state_label,
                mission=st.mission,
                stats={**st.run_stats, **st.mission_stats},
            )
        except Exception:
            pass

        self._auto_tick += 1
        if self._auto_tick >= 5:  # ~1s
            self._auto_tick = 0
            self._run_automations(st)
            self.auto_status.configure(text=self.automation.status_line())

        self.after(200, self._tick_status)

    def _on_auto_changed(self) -> None:
        self._sync_automation_from_ui()
        self._save_config()
        self.auto_status.configure(text=self.automation.status_line())

    def _sync_automation_from_ui(self) -> None:
        s = self.automation.settings
        s.schedule_enabled = bool(self.sched_enabled.get())
        s.bat_start_enabled = bool(self.bat_start_enabled.get())
        s.bat_dock_enabled = bool(self.bat_dock_enabled.get())
        try:
            hh, mm = self.sched_time.get().strip().split(":")
            s.schedule_hour = max(0, min(23, int(hh)))
            s.schedule_minute = max(0, min(59, int(mm)))
            self.sched_time.set(f"{s.schedule_hour:02d}:{s.schedule_minute:02d}")
        except Exception:
            s.schedule_hour, s.schedule_minute = 9, 0
            self.sched_time.set("09:00")
        try:
            s.bat_start_pct = max(1, min(100, int(self.bat_start_pct.get().strip())))
            self.bat_start_pct.set(str(s.bat_start_pct))
        except Exception:
            s.bat_start_pct = 50
            self.bat_start_pct.set("50")
        try:
            s.bat_dock_pct = max(1, min(100, int(self.bat_dock_pct.get().strip())))
            self.bat_dock_pct.set(str(s.bat_dock_pct))
        except Exception:
            s.bat_dock_pct = 15
            self.bat_dock_pct.set("15")
        if s.bat_dock_pct >= s.bat_start_pct and s.bat_start_enabled and s.bat_dock_enabled:
            # Keep dock threshold below start threshold
            s.bat_dock_pct = max(1, s.bat_start_pct - 5)
            self.bat_dock_pct.set(str(s.bat_dock_pct))

    def _apply_automation_to_ui(self) -> None:
        s = self.automation.settings
        self.sched_enabled.set(s.schedule_enabled)
        self.sched_time.set(f"{s.schedule_hour:02d}:{s.schedule_minute:02d}")
        self.bat_start_enabled.set(s.bat_start_enabled)
        self.bat_start_pct.set(str(s.bat_start_pct))
        self.bat_dock_enabled.set(s.bat_dock_enabled)
        self.bat_dock_pct.set(str(s.bat_dock_pct))
        self.auto_status.configure(text=self.automation.status_line())

    def _run_automations(self, st) -> None:
        if self._auto_running or not st.connected or st.connecting:
            return
        self._sync_automation_from_ui()
        action = self.automation.evaluate(
            connected=st.connected,
            battery_pct=st.battery_pct,
            phase=st.phase,
        )
        if not action:
            return

        self._auto_running = True
        self._save_config()  # persist last_schedule_key / armed flags

        def work() -> None:
            msg = ""
            err = ""
            try:
                if action.kind == "clean":
                    msg = self.driver.smart_clean()
                elif action.kind == "dock":
                    msg = self.driver.smart_dock()
                msg = f"{action.reason} → {msg}"
            except Exception as exc:
                err = str(exc)

            def done() -> None:
                self._auto_running = False
                if err:
                    self.error_label.configure(text=err)
                else:
                    self.cmd_label.configure(text=f"Last action: {msg}")
                    self.error_label.configure(text="")
                self.auto_status.configure(text=self.automation.status_line())
                self._save_config()
                # Reflect schedule disable if one-shot
                self._apply_automation_to_ui()

            self.after(0, done)

        threading.Thread(target=work, daemon=True).start()

    def _load_config(self) -> None:
        if not CONFIG_PATH.exists():
            return
        try:
            data = json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
            self.ip_var.set(data.get("ip", ""))
            self.blid_var.set(data.get("blid", ""))
            self.password_var.set(data.get("password", ""))
            if "auto_connect" in data:
                self.auto_var.set(bool(data["auto_connect"]))
            auto_data = data.get("automation") or data
            self.automation.settings = AutomationSettings.from_dict(auto_data)
            # Prefer nested automation block if present
            if "automation" in data:
                self.automation.settings = AutomationSettings.from_dict(data["automation"])
            elif "schedule" in data:
                self.automation.settings = AutomationSettings.from_dict(data)
            self._apply_automation_to_ui()
        except Exception:
            pass

    def _save_config(self) -> None:
        self._sync_automation_from_ui()
        data = {
            "ip": self.ip_var.get().strip(),
            "blid": self.blid_var.get().strip(),
            "password": self.password_var.get().strip(),
            "auto_connect": bool(self.auto_var.get()),
            "automation": self.automation.settings.to_dict(),
        }
        try:
            CONFIG_PATH.write_text(json.dumps(data, indent=2) + "\n", encoding="utf-8")
        except Exception:
            pass

    def _cleanup(self) -> None:
        try:
            self.driver.disconnect()
        except Exception:
            pass

    def on_quit(self) -> None:
        self._closing = True
        try:
            self._save_config()
        except Exception:
            pass
        self._cleanup()
        self.destroy()


def main() -> None:
    app = RoombaApp()
    app.mainloop()


if __name__ == "__main__":
    main()