#!/usr/bin/env python3
"""
vla_bridge.py — real policy forward pass, wired to a robot over zelpi's HAL.

This is the concrete, end-to-end proof of the "install an open VLA/policy,
drive any robot with it over HAL/ROS" story documented in docs/HARDWARE.md:
it loads an actual downloaded checkpoint (see `zelpi hub install <key>`),
builds an observation in exactly the shape that checkpoint expects (from its
own config.json — this is NOT a generic interface, see the honest-limitation
note in the docs), runs one real forward pass through the real model weights,
and sends the resulting action to a running robot-agent via hal_client.py's
"setJoints" skill — the same code path a ROS-driven arm would receive it
through (server/createRosDriver.mjs publishes it to a joint_command topic).

Not hardcoded to SmolVLA: the policy family (smolvla, act, diffusion, pi0,
...) is auto-detected from the checkpoint's own saved config
(`PreTrainedConfig.from_pretrained` -> `cfg.type` -> `get_policy_class`), the
same generic mechanism lerobot's own training/eval scripts use — so any
lerobot-native checkpoint (`zelpi hub install act` / `diffusion-policy` /
`smolvla` / `pi0`, ...) works with this same script.

Since no real camera rig or matching robot is attached, camera inputs are
synthetic (random tensors, correct shape/dtype) and, for language-conditioned
policies, the task string is a static instruction — this proves the plumbing
(checkpoint -> real inference -> HAL -> robot-agent) is genuine, not that the
policy has been fine-tuned for whatever robot receives the action.

Usage:
    python vla_bridge.py --checkpoint F:/.zelpi-models/smolvla --hal-url ws://127.0.0.1:9091
    python vla_bridge.py --checkpoint F:/.zelpi-models/act --hal-url ws://127.0.0.1:9091
"""
from __future__ import annotations

import argparse
import pathlib
import sys
import time

import numpy as np
import torch

sys.path.insert(0, str(pathlib.Path(__file__).parent))
from hal_client import HalClient  # noqa: E402


def build_observation(config):
    """Synthetic observation matching this checkpoint's exact declared shapes."""
    obs = {}
    for key, feat in config.input_features.items():
        if key.startswith("observation.images"):
            obs[key] = torch.rand(*feat.shape, dtype=torch.float32)
        elif key == "observation.state":
            obs[key] = torch.zeros(*feat.shape, dtype=torch.float32)
        else:
            obs[key] = torch.zeros(*feat.shape, dtype=torch.float32)
    if hasattr(config, "vlm_model_name") or getattr(config, "type", "") in ("pi0", "pi05", "smolvla"):
        obs["task"] = "pick up the object"
    return obs


def main():
    ap = argparse.ArgumentParser(description="Run a real policy forward pass and drive a robot over HAL")
    ap.add_argument("--checkpoint", default="F:/.zelpi-models/smolvla")
    ap.add_argument("--hal-url", default="ws://127.0.0.1:9091")
    ap.add_argument("--robot-id", default="vla-bridge-0")
    ap.add_argument("--steps", type=int, default=3, help="number of inference+send steps")
    args = ap.parse_args()

    from lerobot.configs.policies import PreTrainedConfig
    from lerobot.policies.factory import get_policy_class, make_pre_post_processors

    print(f"[vla_bridge] reading config from {args.checkpoint} ...")
    config = PreTrainedConfig.from_pretrained(args.checkpoint)
    policy_cls = get_policy_class(config.type)
    print(f"[vla_bridge] detected policy type: {config.type} -> {policy_cls.__name__}")

    print(f"[vla_bridge] loading checkpoint weights ...")
    policy = policy_cls.from_pretrained(args.checkpoint)
    policy.eval()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cpu":
        print(
            "[vla_bridge] WARNING: no CUDA GPU detected — running on CPU. "
            "Observed forward-pass latency on CPU: ~0.4s (ACT, small) to ~40-60s "
            "(SmolVLA/Diffusion Policy, flow-matching/diffusion denoising). This is "
            "fine for proving the pipeline works end-to-end, but NOT usable for "
            "real-time robot control without a GPU. See docs/HARDWARE.md."
        )
    # The checkpoint's saved config pins whatever device it was trained or
    # migrated on (often "cpu" or "cuda" from a *different* machine) — the
    # preprocessor override below moves the INPUTS to our device, but the
    # policy weights only follow if we move them explicitly. Without this,
    # GPU runs crash with "Expected all tensors to be on the same device".
    policy.to(device)
    device_override = {"device_processor": {"device": device}}
    preprocessor, postprocessor = make_pre_post_processors(
        policy.config,
        pretrained_path=args.checkpoint,
        preprocessor_overrides=device_override,
        postprocessor_overrides=device_override,
    )
    print(f"[vla_bridge] policy loaded: {type(policy).__name__}, action dim {policy.config.output_features['action'].shape}")

    print(f"[vla_bridge] connecting to robot-agent at {args.hal_url} ...")
    client = HalClient(args.hal_url, robot_id=args.robot_id)
    client.connect(timeout=5)
    print(f"[vla_bridge] capabilities: {client.capabilities}")
    client.seed_pose(0, 0, 0)

    for step in range(args.steps):
        obs = build_observation(policy.config)
        t0 = time.monotonic()
        batch = preprocessor(obs)
        with torch.no_grad():
            action = policy.select_action(batch)
        action = postprocessor(action)
        dt = time.monotonic() - t0
        action_np = action.squeeze(0).cpu().numpy() if action.dim() > 1 else action.cpu().numpy()
        joints = [float(x) for x in np.asarray(action_np).flatten()]
        print(f"[vla_bridge] step {step}: real forward pass in {dt:.2f}s -> action {joints}")

        result = client.send_joints(joints, wait=True, timeout=3.0)
        print(f"[vla_bridge] setJoints result: {result}")
        time.sleep(0.2)
        print(f"[vla_bridge] telemetry after setJoints: {client.telemetry}")

    client.close()
    print("[vla_bridge] done")


if __name__ == "__main__":
    main()