Bolt-Assistant/bolt_assistant/heartbeat.py

"""The heartbeat that brings Bolt's reminders and delivered files in.

Same protocol as Bolt-Pet's heartbeat (desk_client/bolt_desk.py):

    POST /desk/report_status  {"session_id", "download_dir"}
        -> {"ok": true} or {"reply": "<announcement to speak>"} (plus
           "suggestions" when there are any)
    GET  /desk/files?session_id=...     -> {"files": [{"id", "name", "size"}]}
    GET  /desk/files/<id>?session_id=...-> the bytes; the server *dequeues*
           on this GET, so write to a .part first and rename only once the
           whole file is here, and never overwrite an existing file.

A heartbeat thread is the only thing standing between a delivered file and
this machine, so every failure is logged and swallowed: one bad beat must
never stop the loop, and a refused key gets a loud warning once at startup
instead of silent nothing-ever-arrives.

Announcements returned by report_status are popped server-side the moment
they are handed out, so the caller (__main__) feeds them to
assistant.announce() rather than dropping them. Everything network-shaped
is injectable for the tests.
"""

from __future__ import annotations

import os
import re
import time
from pathlib import Path
from typing import Callable, List, Optional

import requests

from . import config
from .server import _headers

log: Callable[[str], None] = print


def _url(path: str) -> str:
    return f"{config.SERVER_URL}{path}"


def _safe_filename(name: str) -> str:
    """A server-sent name is a display hint, not a trusted path: strip path
    separators, control characters and leading dots before it touches disk."""
    cleaned = re.sub(r'[\\/:*?"<>|\x00-\x1f]', "_", str(name or "")).strip().strip(".")
    return cleaned or "download"


def _free_path(folder: Path, name: str) -> Path:
    """A name that doesn't exist yet: the server never overwrites, and neither
    do we. "note (2).txt" after "note.txt" and "note (1).txt", like a browser."""
    target = folder / name
    if not target.exists():
        return target
    stem, suffix = target.stem, target.suffix
    for n in range(1, 1000):
        candidate = folder / f"{stem} ({n}){suffix}"
        if not candidate.exists():
            return candidate
    return folder / f"{stem} ({int(time.time())}){suffix}"


def fetch_server_health(timeout: float = 10.0, get: Callable = requests.get) -> None:
    """One keyed probe at startup. A refusal gets a loud warning — a dead
    DESK_API_KEY means total silence (no answers, no files, no reminders)
    with no other symptom."""
    try:
        response = get(_url("/desk/health"), headers=_headers(), timeout=timeout)
        info = response.json()
        log(f"server: {info}")
        if info.get("key_ok") is False:
            log("!! the server refused this device's DESK_API_KEY — no files, no reminders,")
            log("!! and no answers will reach this assistant until it is updated")
            log("!! set DESK_API_KEY in .env — it is re-read on the next request,")
            log("!! so no restart is needed once the file is corrected")
    except Exception as exc:
        log(f"warning: server not reachable yet ({exc}) — will keep trying per-request")


def _download(file_id: str, name: str, folder: Path,
              get: Callable = requests.get) -> Optional[Path]:
    """Fetch one queued file and save it. The server dequeues on this GET, so
    write to a .part first and rename only once the bytes are all here."""
    try:
        response = get(
            _url(f"/desk/files/{file_id}"),
            params={"session_id": config.SESSION_ID}, headers=_headers(),
            timeout=(10, 120), stream=True,
        )
        if response.status_code != 200:
            log(f"file download refused ({response.status_code}) for {name}")
            return None
        folder.mkdir(parents=True, exist_ok=True)
        target = _free_path(folder, _safe_filename(name))
        part = target.with_name(target.name + ".part")
        try:
            with open(part, "wb") as handle:
                for chunk in response.iter_content(65536):
                    if chunk:
                        handle.write(chunk)
            os.replace(part, target)
        except Exception:
            try:
                part.unlink()
            except OSError:
                pass
            raise
        return target
    except Exception as exc:
        log(f"file download failed for {name}: {exc}")
        return None


def _fetch_outbox(folder: Path, get: Callable = requests.get) -> List[Path]:
    try:
        response = get(
            _url("/desk/files"),
            params={"session_id": config.SESSION_ID}, headers=_headers(),
            timeout=(5, 30),
        )
    except Exception as exc:
        log(f"couldn't list queued files: {exc}")
        return []
    if response.status_code != 200:
        return []
    try:
        body = response.json()
    except ValueError:
        return []
    entries = body.get("files") if isinstance(body, dict) else None
    saved: List[Path] = []
    for entry in entries or []:
        if not isinstance(entry, dict) or not entry.get("id"):
            continue
        target = _download(str(entry["id"]), str(entry.get("name") or "download"), folder, get)
        if target is not None:
            saved.append(target)
    return saved


def heartbeat_once(download_dir=None, post: Callable = requests.post,
                   get: Callable = requests.get) -> dict:
    """One beat: report in (reminders, files and presence all flow), then
    collect whatever the server is holding for this device. Returns
    {"ok", "announcement", "saved"} — *announcement* is something Bolt wants
    spoken, *saved* the Paths just downloaded."""
    folder = Path(download_dir) if download_dir is not None else config.DOWNLOAD_DIR
    result: dict = {"ok": False, "announcement": "", "saved": []}
    try:
        folder.mkdir(parents=True, exist_ok=True)
    except OSError as exc:
        log(f"can't create the download folder {folder}: {exc}")
    payload = {"session_id": config.SESSION_ID, "download_dir": str(folder)}
    try:
        response = post(_url("/desk/report_status"), json=payload,
                        headers=_headers(), timeout=20)
    except Exception as exc:
        log(f"heartbeat failed: {exc}")
        return result
    if response.status_code != 200:
        return result
    try:
        body = response.json()
    except ValueError:
        body = {}
    announcement = body.get("reply") if isinstance(body, dict) else None
    if isinstance(announcement, str) and announcement.strip():
        result["announcement"] = " ".join(announcement.split())
    result["saved"] = _fetch_outbox(folder, get)
    if result["saved"]:
        names = ", ".join(path.name for path in result["saved"])
        log(f"saved {len(result['saved'])} file(s) to {folder}: {names}")
    result["ok"] = True
    return result


def heartbeat_loop(stop, on_result: Optional[Callable[[dict], None]] = None) -> None:
    """Beat every HEARTBEAT_SECONDS until *stop* is set, handing each result
    to *on_result*. A bad iteration is logged, never fatal; a typo in the
    interval must not become a request flood (10 s floor)."""
    if config.HEARTBEAT_SECONDS <= 0:
        return
    interval = max(10, config.HEARTBEAT_SECONDS)
    while not stop.is_set():
        try:
            result = heartbeat_once()
            if on_result is not None:
                on_result(result)
        except Exception as exc:
            log(f"heartbeat iteration failed: {exc}")
        stop.wait(interval)