modules/ar_scene.py

from __future__ import annotations

"""
CANONICAL AR ANCHOR STORE — AR-FIX-SESSION
This module is the single source of truth for AR scene state across all clients.
All anchor writes from any client (WebAR, backend ANCHOR_DB, Tauri SQLite) should
route through apply_patch() to ensure cross-client visibility via the WebSocket bus.
State is persisted to data/ar_scene.json and broadcast via phone/ws_scene_bus.py.

modules/ar_scene.py - Shared AR "scene state" for FRIDAY ecosystem (PC <-> Phone).

This is intentionally lightweight:
- Authoritative state lives on the PC (this repo runtime).
- Phone and PC exchange small JSON patches (no binary/video transport).
- Rendering / tracking modules can subscribe to state changes.

We start with state for:
- active preset (table hologram / helmet HUD / free float)
- anchor poses (4x4 matrices) for ArUco or face pose
- panel placement overrides for AR overlay
"""

import json
import os
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
DATA_DIR = ROOT / "data"
DATA_DIR.mkdir(parents=True, exist_ok=True)
STATE_PATH = DATA_DIR / "ar_scene.json"

# S4: RLock, not Lock — apply_patch holds the lock and then calls load_state(),
# which acquires it again. With a plain Lock every patch deadlocked the calling
# thread (and froze the scene bus event loop) on the nested acquire.
_lock = threading.RLock()
_state: dict[str, Any] | None = None


def _now_ms() -> int:
    return int(time.time() * 1000)


def _default_state() -> dict[str, Any]:
    return {
        "v": 1,
        "updated_ms": _now_ms(),
        "preset": "table",  # table | helmet | float
        "enabled": False,
        "ar_objects": [
            # authoritative scene objects visible to every client (phone AR Lab,
            # WebAR, headsets). Each: {id,type,label,px,py,pz[,model_url]}.
            # This is the exact key the mobile scene client reads on hello/state.
        ],
        "anchors": {
            # key -> 4x4 row-major float matrix
            # example: {"aruco:23": [ ... 16 floats ... ]}
        },
        "ui": {
            # optional AR overlay panel layout overrides
            "panels": {},
        },
    }


def load_state() -> dict[str, Any]:
    global _state
    with _lock:
        if _state is not None:
            return dict(_state)
        try:
            if STATE_PATH.exists():
                raw = json.loads(STATE_PATH.read_text(encoding="utf-8"))
                if isinstance(raw, dict) and raw.get("v") == 1:
                    _state = raw
                else:
                    _state = _default_state()
            else:
                _state = _default_state()
        except Exception:
            _state = _default_state()
        return dict(_state)


def save_state(state: dict[str, Any]) -> None:
    try:
        tmp = str(STATE_PATH) + ".tmp"
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(state, f, indent=2)
        os.replace(tmp, STATE_PATH)
    except Exception:
        pass


@dataclass(frozen=True)
class PatchResult:
    ok: bool
    state: dict[str, Any]
    message: str = ""


def apply_patch(patch: dict[str, Any]) -> PatchResult:
    """
    Apply a shallow patch to the AR scene state.

    Supported patch keys:
      - enabled: bool
      - preset: "table"|"helmet"|"float"
      - anchors: dict[str, list[float]] (replaces/merges per key)
      - ui: dict (merged shallow)
      - ar_objects: list[dict] (full-list replace — REST /ar_scene sync)
      - action + objectId: SPAWN|MANIPULATE|DELETE object verb (phone scene WS)
      - action "spawn-generated-glb" + model_url/model_name (AI model push)

    Returns the updated full state.
    """
    if not isinstance(patch, dict):
        return PatchResult(ok=False, state=load_state(), message="bad patch")

    global _state
    with _lock:
        st = load_state()
        try:
            if "enabled" in patch:
                st["enabled"] = bool(patch.get("enabled"))
            if "preset" in patch:
                p = str(patch.get("preset") or "").strip().lower()
                if p in ("table", "helmet", "float"):
                    st["preset"] = p
            if "anchors" in patch and isinstance(patch.get("anchors"), dict):
                a = st.get("anchors")
                if not isinstance(a, dict):
                    a = {}
                for k, v in (patch.get("anchors") or {}).items():
                    kk = str(k)[:64]
                    if isinstance(v, list) and len(v) == 16:
                        a[kk] = [float(x) for x in v]
                st["anchors"] = a
            if "ui" in patch and isinstance(patch.get("ui"), dict):
                ui = st.get("ui")
                if not isinstance(ui, dict):
                    ui = {}
                for k, v in (patch.get("ui") or {}).items():
                    ui[str(k)[:64]] = v
                st["ui"] = ui

            # S4 (AR overhaul): object-level scene sync. Both the phone AR Lab and
            # the generated-GLB model push need authoritative object state under
            # `ar_objects` — the exact key the mobile scene client reads back on
            # hello/state. Previously apply_patch ignored all of this, so spawns,
            # manipulations, room models and pushed models never became
            # authoritative and never synced across clients.
            objects = st.get("ar_objects")
            if not isinstance(objects, list):
                objects = []

            def _clean_obj(raw: dict[str, Any]) -> dict[str, Any] | None:
                oid = str(raw.get("id") or raw.get("objectId") or "").strip()
                if not oid:
                    return None
                pz = raw.get("pz")
                obj: dict[str, Any] = {
                    "id": oid[:64],
                    "type": str(raw.get("type") or "CUBE")[:64],
                    "label": str(raw.get("label") or "Object")[:120],
                    "px": float(raw.get("px") or 0.0),
                    "py": float(raw.get("py") or 0.0),
                    "pz": float(pz) if pz is not None else -2.0,
                }
                murl = raw.get("model_url")
                if murl:
                    obj["model_url"] = str(murl)[:512]
                return obj

            # Full-list replace (REST /ar_scene full-scene sync from the phone).
            if isinstance(patch.get("ar_objects"), list):
                rebuilt: list[dict[str, Any]] = []
                for raw in patch["ar_objects"]:
                    if isinstance(raw, dict):
                        c = _clean_obj(raw)
                        if c is not None:
                            rebuilt.append(c)
                objects = rebuilt

            # Verb patches: phone scene WS ScenePatch + AI model push.
            action = str(patch.get("action") or "").strip()
            if action:
                au = action.upper()
                oid = str(patch.get("objectId") or patch.get("id") or "").strip()[:64]
                if "GENERATED-GLB" in au or au.startswith("SPAWN-GENERATED"):
                    murl = str(patch.get("model_url") or "")[:512]
                    gid = oid or ("glb-" + str(abs(hash(murl)) % 100000))
                    objects = [o for o in objects if o.get("id") != gid]
                    objects.append({
                        "id": gid,
                        "type": "GENERATED_GLB",
                        "label": str(patch.get("model_name") or "Generated Model")[:120],
                        "px": 0.0, "py": 0.0, "pz": -2.0,
                        "model_url": murl,
                    })
                elif au == "SPAWN" and oid:
                    c = _clean_obj(patch)
                    if c is not None:
                        objects = [o for o in objects if o.get("id") != oid] + [c]
                elif au == "MANIPULATE" and oid:
                    updated: list[dict[str, Any]] = []
                    for o in objects:
                        if o.get("id") == oid:
                            o = dict(o)
                            for k in ("px", "py", "pz"):
                                if patch.get(k) is not None:
                                    o[k] = float(patch.get(k))
                            if patch.get("type"):
                                o["type"] = str(patch.get("type"))[:64]
                            if patch.get("label"):
                                o["label"] = str(patch.get("label"))[:120]
                        updated.append(o)
                    objects = updated
                elif au == "DELETE" and oid:
                    objects = [o for o in objects if o.get("id") != oid]

            st["ar_objects"] = objects[-200:]

            st["updated_ms"] = _now_ms()
            _state = st
            save_state(st)
            return PatchResult(ok=True, state=dict(st), message="ok")
        except Exception as e:
            return PatchResult(ok=False, state=dict(st), message=str(e)[:160])