backend/routes/xr_routes.py

from fastapi import HTTPException, APIRouter, Request, Depends
from pydantic import BaseModel

# Server-side plan gating + per-user usage metering for the company-provisioned
# 3D-model-generation routes (spawn/builder/forge). AR_MODEL_GEN is only granted on
# paid plans and is metered against the plan's cap; the frontend renders the locked
# state but enforcement lives here.
from backend.billing.gating import gate_and_meter
from backend.billing.plans import Feature
_AR_MODEL_GEN_GATE = Depends(gate_and_meter(Feature.AR_MODEL_GEN))

# AR-FIX-SESSION: import canonical AR scene store to mirror anchor writes.
# Guarded by try/except — safe in frozen PyInstaller context if path not set up.
try:
    from modules.ar_scene import apply_patch as _ar_scene_apply_patch
except ImportError:
    _ar_scene_apply_patch = None

router = APIRouter()

class SpatialAnchor(BaseModel):
    id: str
    x: float
    y: float
    z: float
    session_id: str = "default"

ANCHOR_DB = {}

@router.post("/anchor")
async def save_anchor(a: SpatialAnchor):
    if a.session_id not in ANCHOR_DB:
        ANCHOR_DB[a.session_id] = []

    # Update or insert into ephemeral in-memory store
    status = "saved"
    for i, existing in enumerate(ANCHOR_DB[a.session_id]):
        if existing.id == a.id:
            ANCHOR_DB[a.session_id][i] = a
            status = "updated"
            break
    else:
        ANCHOR_DB[a.session_id].append(a)

    # AR-FIX-SESSION: mirror anchor to canonical ar_scene.json via apply_patch.
    # Constructs 4x4 row-major translation matrix (identity rotation) from x/y/z.
    if _ar_scene_apply_patch is not None:
        _ar_scene_apply_patch({
            "anchors": {
                a.id: [1.0, 0.0, 0.0, a.x,
                       0.0, 1.0, 0.0, a.y,
                       0.0, 0.0, 1.0, a.z,
                       0.0, 0.0, 0.0, 1.0]
            }
        })

    return {"status": status}

@router.get("/anchors/{session_id}")
async def get_anchors(session_id: str):
    return ANCHOR_DB.get(session_id, [])

@router.get("/session_state")
async def get_session_state():
    # AR-FIX-S2: was unconditionally {state: "active"} — now reflects real WS client count.
    # ws_manager.active_connections is a list[WebSocket] maintained by ConnectionManager.connect()
    # and ConnectionManager.disconnect() in backend/ws/agent_ws.py.
    # Import inside function body — same pattern as /ar_task, /ar_scene_patch, etc.
    from backend.ws.agent_ws import ws_manager
    connected = len(ws_manager.active_connections)
    total_anchors = sum(len(v) for v in ANCHOR_DB.values())
    return {
        "state": "active" if connected > 0 else "idle",
        "connected_clients": connected,
        "total_anchors": total_anchors,
    }


class ShaderPayload(BaseModel):
    shader_code: str
    target: str = "webgl2"

@router.post("/compile_shader")
async def compile_shader(payload: ShaderPayload):
    """
    Validate a WebGL/WebXR shader server-side and return a content hash.

    There is no server-side binary compilation here by design: the AR runtime is
    WebGL/Three.js, so the shader is actually compiled client-side by the GPU
    driver at runtime. This endpoint therefore does real, honest work — a
    lightweight static validation pass + a stable content hash the client can use
    for caching — rather than returning a fake `mock://` blob URL (the previous
    behaviour), which read like a real artifact link but resolved to nothing.
    """
    import hashlib
    code = payload.shader_code or ""
    shader_hash = hashlib.md5(code.encode()).hexdigest()

    warnings, errors = [], []
    if not code.strip():
        errors.append("empty shader source")
    if code.count("{") != code.count("}"):
        errors.append("unbalanced braces")
    if code.count("(") != code.count(")"):
        errors.append("unbalanced parentheses")
    if "void main" not in code and code.strip():
        warnings.append("no main() entry point found")

    return {
        "status": "validated" if not errors else "error",
        "hash": shader_hash,
        "target": payload.target,
        "warnings": warnings,
        "errors": errors,
        # Compilation is client-side (WebGL); stated honestly, not a fake URL.
        "compilation": "client_side_webgl",
    }

class ArTaskPayload(BaseModel):
    task: str
    mode: str = "task"
    device_name: str = "unknown"

@router.post("/ar_task")
async def handle_ar_task(payload: ArTaskPayload):
    """
    Receives AR tasks from Android phone or WebAR frontend,
    and broadcasts them to all connected XR clients (e.g. headsets, Desktop UI).
    """
    from backend.ws.agent_ws import ws_manager
    await ws_manager.broadcast({
        "event": "ar:task",
        "payload": payload.dict()
    })
    return {"ok": True, "status": "dispatched", "message": f"Task '{payload.task}' sent to AR bus."}

@router.get("/ar_config")
async def get_ar_config(request: Request):
    """
    Returns base configuration for the WebAR engine.

    Everything here is derived from the *actual request*, never hardcoded. The
    previous version returned ``ws://{HOST or 127.0.0.1}:5050`` and
    ``http://127.0.0.1:7474/...`` unconditionally, which is wrong for every
    client that is not on the same machine as the backend: a phone or headset
    resolves 127.0.0.1 to itself, so the scene socket and game URLs could never
    connect from the cloud deployment. Port 5050 was wrong too — the socket is
    mounted at ``/scene-ws`` on this same app (backend/main.py), and a hosted
    Space only exposes 443.
    """
    # Trust X-Forwarded-Proto before request.url.scheme. HuggingFace terminates
    # TLS at its proxy and forwards plain HTTP to the container, so base_url
    # reports "http" even though the browser loaded the page over https. Deriving
    # the socket scheme from it produced `ws://…` on an `https://` page, which
    # browsers block outright as mixed content — trading the old 127.0.0.1 bug
    # for a subtler one that only shows up in a real browser.
    forwarded = request.headers.get("x-forwarded-proto", "").split(",")[0].strip().lower()
    scheme = forwarded or (request.url.scheme or "http")
    host = str(request.base_url).split("://", 1)[1].rstrip("/")
    base = f"{scheme}://{host}"
    ws_scheme = "wss" if scheme == "https" else "ws"
    return {
        "ok": True,
        "deviceName": "OMEGA-CORE",
        "wsBase": f"{ws_scheme}://{host}/scene-ws",
        "gameUrls": {
            "eternum": f"{base}/games/eternum",
            "ripples": f"{base}/games/ripples",
        },
    }

@router.post("/ar_scene_patch")
async def handle_ar_scene_patch(patch: dict):
    """
    Syncs WebAR component states across the network.
    """
    from backend.ws.agent_ws import ws_manager
    await ws_manager.broadcast({
        "event": "ar:patch",
        "payload": patch
    })
    return {"ok": True, "status": "synced"}

@router.get("/ar_scene")
async def get_ar_scene_state():
    # Reflect the authoritative object list so a fresh client can hydrate over
    # REST (the WS scene bus also delivers it via hello/state).
    if _ar_scene_apply_patch is not None:
        try:
            from modules.ar_scene import load_state
            st = load_state()
            return {
                "ok": True,
                "entities": st.get("ar_objects", []),
                "environment": st.get("preset", "default"),
            }
        except Exception:
            pass
    return {"ok": True, "entities": [], "environment": "default"}


class ArSceneSyncPayload(BaseModel):
    device_name: str = "unknown"
    objects: list[dict] = []


@router.post("/ar_scene")
async def sync_ar_scene(payload: ArSceneSyncPayload):
    """
    Full-scene sync from the mobile AR Lab. Guardian POSTs {device_name, objects}
    on every spawn / manipulate / delete / room-scan. Previously the backend only
    exposed GET /ar_scene, so this POST returned 405 on cloud and the authoritative
    scene never persisted.

    Persists the object list into the canonical store and fans it out to every
    connected scene-bus client (phone reads state.ar_objects) plus the agent WS
    event bus (desktop / headset listeners), mirroring the model-push path.
    """
    objects = payload.objects if isinstance(payload.objects, list) else []

    # Durable authoritative write (idempotent full-list replace).
    if _ar_scene_apply_patch is not None:
        _ar_scene_apply_patch({"ar_objects": objects})

    # Fan out to scene-bus clients — the bus re-applies + broadcasts {type:state}.
    try:
        from phone import ws_scene_bus
        await ws_scene_bus.broadcast_to_clients_async({
            "type": "patch",
            "patch": {"ar_objects": objects},
        })
    except Exception:
        pass

    # Parity with the other AR handlers: notify the agent WS event bus too.
    try:
        from backend.ws.agent_ws import ws_manager
        await ws_manager.broadcast({
            "event": "ar:scene",
            "payload": {"device_name": payload.device_name, "objects": objects},
        })
    except Exception:
        pass

    return {"ok": True, "status": "synced", "count": len(objects)}

@router.post("/ar_game_state")
async def update_ar_game_state(state: dict):
    from backend.ws.agent_ws import ws_manager
    await ws_manager.broadcast({"event": "ar:game_state", "payload": state})
    return {"ok": True}

@router.post("/ar_game_builder")
async def sync_ar_game_builder(graph: dict):
    from backend.ws.agent_ws import ws_manager
    await ws_manager.broadcast({"event": "ar:game_builder", "payload": graph})
    return {"ok": True}

class SpawnModelRequest(BaseModel):
    description: str
    persona: str = "JARVIS"
    # Without this the endpoint was a dead end: spawn_3d_model() returns an
    # "ask the user, then call again with use_cloud_fallback=True" instruction
    # whenever local SF3D weights are absent — which is ALWAYS true on the
    # CPU-only cloud Space — but the request model had no way to express it,
    # so no client could ever take the cloud path.
    use_cloud_fallback: bool = False

@router.post("/model3d/spawn")
async def spawn_model(request: SpawnModelRequest, _user: str = _AR_MODEL_GEN_GATE):
    import logging
    logger = logging.getLogger(__name__)
    from backend.tools import xr_tools
    try:
        # On the cloud Space the 10GB local weights are never present, so honour
        # an explicit cloud-fallback request instead of returning the interactive
        # prompt string as if it were a model path.
        model_path = await xr_tools.spawn_3d_model(
            request.description, request.persona, request.use_cloud_fallback
        )
        if isinstance(model_path, str) and model_path.startswith("STATUS: LOCAL_MODEL_MISSING"):
            # Surface this as a real, actionable API state — not a fake success
            # whose "model_path" is a paragraph of instructions.
            return {
                "status": "needs_confirmation",
                "reason": "local_sf3d_missing",
                "message": "Local Stable-Fast-3D weights are not installed. Retry with "
                           "use_cloud_fallback=true to generate via the hosted API.",
            }
        return {"status": "ok", "model_path": model_path}
    except xr_tools.StableFast3DNotProvisionedError as e:
        logger.error(f"SF3D not provisioned: {e}")
        raise HTTPException(status_code=503, detail=str(e))
    except Exception as e:
        # A core feature must not hard-fail because a THIRD-PARTY space is down.
        # Verified 2026-07-20: stabilityai/stable-fast-3d's /run_button raises an
        # opaque AppError for every parameter combination, reproduced directly
        # from a client outside this backend — their GPU inference is the fault,
        # not our request (signature confirmed against their own view_api).
        # Our procedural Blender Builder is fully operational, so degrade to it
        # instead of returning a 500: the user still gets a model in the scene.
        logger.warning(f"SF3D generation failed ({e}); falling back to procedural Builder.")
        try:
            result = await builder_generate(
                BuilderRequest(description=request.description, persona=request.persona)
            )
            if isinstance(result, dict):
                result = dict(result)
                result["fallback"] = "procedural_builder"
                result["fallback_reason"] = f"stable-fast-3d unavailable: {str(e)[:160]}"
            return result
        except HTTPException:
            raise
        except Exception as fallback_error:
            logger.exception("Procedural Builder fallback also failed")
            raise HTTPException(
                status_code=502,
                detail=f"SF3D unavailable ({str(e)[:120]}) and Builder fallback failed: {fallback_error}",
            )

@router.get("/model3d/health")
async def check_health():
    from backend.services.connectors import stable_fast_3d_local
    return await stable_fast_3d_local.health_check()


# ─── S4: Builder + Forge Blender pipeline ────────────────────────────────────

BLUEPRINT_PROMPT = """Create a compact JSON blueprint for a fully interactive 3D model.
Return JSON only, no markdown, no prose.
Schema:
{{
  "name": "string",
  "kit": "house|room|device|vehicle|character|scene",
  "parts": [
    {{
      "name": "string",
      "type": "base|wall|door|window|roof|drawer|panel|screen|wheel|prop|light|arm|leg|head|detail",
      "geometry": "box|sphere|cylinder|cone|torus|ring|capsule|octahedron|icosahedron",
      "material": "matte|metal|glass|glow|wood",
      "position": [0,0,0],
      "rotation": [0,0,0],
      "scale": [1,1,1],
      "actions": ["open","close","toggle","rotate","reveal","hide","shatter","reconfigure"]
    }}
  ],
  "behaviors": [
    {{ "gesture": "pinch|fist|open-palm|swipe-left|swipe-right|two-hand-scale", "action": "string", "target": "part name" }}
  ]
}}
Prompt: {description}"""


def _parse_blueprint_json(text: str) -> dict | None:
    """Extract the blueprint object from LLM output.

    S4: the NVIDIA fallback (glm-5.1 and friends) wraps JSON in <think> blocks
    and prose, and may emit several JSON-ish fragments. A first-{ to last-}
    slice fails whenever any prose contains a stray brace — scan for the first
    *balanced* object that parses and looks like a blueprint instead.
    """
    import json as _json
    import re as _re
    raw = str(text or "")
    raw = _re.sub(r"<think>.*?</think>", "", raw, flags=_re.DOTALL)
    raw = raw.replace("```json", "```").strip()
    candidates = []
    # Fenced blocks first — models that follow instructions put the JSON there.
    for block in raw.split("```")[1::2]:
        candidates.append(block.strip())
    candidates.append(raw.replace("```", ""))
    for candidate in candidates:
        idx = candidate.find("{")
        while idx >= 0:
            depth = 0
            for end in range(idx, len(candidate)):
                ch = candidate[end]
                if ch == "{":
                    depth += 1
                elif ch == "}":
                    depth -= 1
                    if depth == 0:
                        try:
                            obj = _json.loads(candidate[idx:end + 1])
                        except Exception:
                            obj = None
                        if isinstance(obj, dict) and isinstance(obj.get("parts"), list) and obj["parts"]:
                            return obj
                        break
            idx = candidate.find("{", idx + 1)
    return None


class BuilderRequest(BaseModel):
    description: str
    persona: str = "JARVIS"
    name: str = ""


@router.post("/model3d/builder")
async def builder_generate(request: BuilderRequest, _user: str = _AR_MODEL_GEN_GATE):
    """"Builder <description>" — blueprint via the token manager (it selects the
    right vault credential for the generation task, NVIDIA fallback included),
    then a real headless-Blender procedural build, then serve + durable hub push."""
    import logging
    logger = logging.getLogger(__name__)
    from backend.services import blender_processor
    from backend.services.token_manager import gemini_call_with_checkpoint

    if not blender_processor.blender_available():
        raise HTTPException(status_code=503, detail="Blender is not installed in this runtime.")

    description = request.description.strip()
    if not description:
        raise HTTPException(status_code=422, detail="A model description is required.")

    # Token manager owns credential selection: task_type="image" routes to the
    # IMAGE_GENERATION_* key domain (cloud/pc aware), with NVIDIA vault fallback.
    raw = await gemini_call_with_checkpoint(
        BLUEPRINT_PROMPT.format(description=description),
        task_type="image", persona=request.persona.lower(),
    )
    blueprint = _parse_blueprint_json(raw)
    if not blueprint:
        head = str(raw or "")[:500]
        logger.error(f"Builder blueprint unparseable; raw head: {head!r}")
        raise HTTPException(
            status_code=502,
            detail=f"Blueprint generation returned no parseable JSON. Model said: {head[:160]!r}",
        )

    name = request.name or blueprint.get("name") or description[:24]
    try:
        result = await blender_processor.build_from_blueprint(blueprint, name)
    except Exception as e:
        logger.exception("Builder Blender pass failed")
        raise HTTPException(status_code=500, detail=f"Blender build failed: {e}")

    from backend.tools.xr_tools import push_model_to_ar_scene
    await push_model_to_ar_scene(result["model_path"], name)
    return {"status": "ok", "name": name, "blueprint": blueprint, **result}


class ForgeProcessRequest(BaseModel):
    model_url: str
    name: str = "Sourced Model"
    kit: str = "scene"
    behaviors: list | None = None


@router.post("/model3d/forge_process")
async def forge_process(request: ForgeProcessRequest, _user: str = _AR_MODEL_GEN_GATE):
    """Forge handoff for sourced models (Sketchfab/TurboSquid/...). Reuse-check runs
    FIRST against saved templates (pure-python GLB part scan — no Blender launch);
    Blender only on a miss, and the fresh result becomes a new reusable template."""
    import logging
    import os
    import tempfile
    logger = logging.getLogger(__name__)
    from backend.services import blender_processor

    url = request.model_url.strip()
    if not url.lower().startswith(("http://", "https://")):
        raise HTTPException(status_code=422, detail="model_url must be an http(s) URL.")

    import asyncio as _asyncio
    import requests as _requests

    def _download() -> str:
        with _requests.get(url, stream=True, timeout=60) as resp:
            resp.raise_for_status()
            fd, path = tempfile.mkstemp(suffix=".glb")
            written = 0
            with os.fdopen(fd, "wb") as fh:
                for chunk in resp.iter_content(1024 * 256):
                    written += len(chunk)
                    if written > 200 * 1024 * 1024:
                        raise RuntimeError("Model exceeds the 200MB processing limit.")
                    fh.write(chunk)
            return path

    try:
        local_path = await _asyncio.to_thread(_download)
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"Model download failed: {e}")

    template = {"behaviors": request.behaviors} if request.behaviors else None
    try:
        result = await blender_processor.forge_process_glb(local_path, request.name, request.kit, template)
    except Exception as e:
        logger.exception("Forge Blender pass failed")
        raise HTTPException(status_code=500, detail=f"Forge processing failed: {e}")
    finally:
        # The temp download only feeds the part scan / Blender import; the processed
        # output is a separate file in storage/models3d. Always clean up.
        if os.path.exists(local_path):
            os.unlink(local_path)

    if result.get("reused"):
        # model_path was the (now deleted) temp file — the client keeps its own copy
        # of the sourced model and only needs the template to apply.
        result.pop("model_path", None)

    return {"status": "ok", **result}


@router.get("/model3d/blender_health")
async def blender_health():
    from backend.services import blender_processor
    return {
        "blender_available": blender_processor.blender_available(),
        "blender_path": blender_processor.find_blender() or None,
        "templates": len(blender_processor.load_templates()),
        "models_dataset": blender_processor.MODELS_DATASET,
    }