#!/usr/bin/env python3
"""Controller capabilities the viewer drives that lerobot does not advertise.

Everything here works against an unmodified ``lerobot`` from main. The demo
pins upstream rather than carrying patches into it, so anything the UI needs
beyond the ``RobotController`` protocol is adapted from the outside.
"""

from __future__ import annotations

import ast
import importlib
import logging
from pathlib import Path

from lerobot.robots.unitree_g1 import unitree_g1

logger = logging.getLogger(__name__)

# Scalar commands a controller keeps on an instance attribute, as
# {controller class: {action key: (min, max, default, attribute)}}.
#
# The protocol has no vocabulary for these: run_step is handed the joystick
# axes, and anything else the policy takes as an observation is reachable only
# as an attribute. The driver applies them by assignment, hence the name.
COMMANDS: dict[str, dict[str, tuple[float, float, float, str]]] = {
    # Commanded base height, i.e. how deep GR00T squats. The policy reads it as
    # an observation and holds whatever it is given rather than tracking to it,
    # so a slider maps straight onto it. Upstream only ever nudges it by
    # 0.001/tick from held gamepad buttons, which no browser is going to send.
    "GrootLocomotionController": {"command.height": (0.50, 1.00, 0.74, "groot_height_cmd")},
}


def limit_inference_threads(threads: int = 1) -> None:
    """Pin ONNX Runtime to ``threads`` per pool, for every session made later.

    Thread pools are sized from the machine's core count, and inside a container
    that is the host's count rather than the CPU quota. On a 2-vCPU Space that
    meant ~60 threads per process contending for 2 cores, which costs far more
    than the inference itself: physics needs about 5% of a core here, yet the
    container sat pegged at its limit and the policies missed their deadlines.

    These are small MLPs, so one thread each beats any amount of contention.
    lerobot builds its sessions as ``ort.InferenceSession(path)`` with no
    options, so the only way in from outside is to wrap the constructor.

    The matching BLAS and OpenMP limits have to be environment variables set
    before numpy and torch are imported, so they live in the Dockerfile.
    """
    import onnxruntime as ort

    if getattr(ort.InferenceSession, "_thread_capped", False):
        return
    original = ort.InferenceSession

    def capped(*args, **kwargs):
        # sess_options is the second positional parameter; only inject when the
        # caller passed none, so an explicit choice is never overridden.
        if len(args) < 2 and kwargs.get("sess_options") is None:
            options = ort.SessionOptions()
            options.intra_op_num_threads = threads
            options.inter_op_num_threads = threads
            options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
            kwargs["sess_options"] = options
        return original(*args, **kwargs)

    capped._thread_capped = True
    ort.InferenceSession = capped
    logger.info(f"onnxruntime capped to {threads} thread(s) per session")


def _base_names(node: ast.ClassDef) -> set[str]:
    """Base classes of a ClassDef as written (``X`` or ``mod.X`` -> ``X``)."""
    return {b.id for b in node.bases if isinstance(b, ast.Name)} | {
        b.attr for b in node.bases if isinstance(b, ast.Attribute)
    }


def discover_controllers() -> dict[str, str]:
    """Map controller class name to module path, from the installed lerobot.

    Anything in ``controllers/`` declaring a ``RobotController`` subclass is
    selectable, so dropping a module in there is enough; nothing here needs
    editing. Deliberately uncached, so a new one appears on a page refresh
    rather than a restart.

    The modules are parsed rather than imported: ``controllers/__init__.py``
    imports all of them eagerly, which pulls in onnxruntime and every policy's
    dependencies just to answer "what can I pick?", and would let a single
    controller with an unmet optional dependency hide all the others.
    """
    package = f"{unitree_g1.__package__}.controllers"
    found: dict[str, str] = {}
    for path in sorted((Path(unitree_g1.__file__).parent / "controllers").glob("[!_]*.py")):
        try:
            tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
        except (OSError, SyntaxError) as e:
            logger.warning(f"skipping controller module {path.name}: {e}")
            continue
        for n in tree.body:
            if isinstance(n, ast.ClassDef) and "RobotController" in _base_names(n):
                found[n.name] = f"{package}.{path.stem}"
    return found


def available_controllers() -> list[str]:
    """Controller class names that can be selected."""
    return sorted(discover_controllers())


def make_controller(name: str):
    """Instantiate a discovered controller by class name.

    Mirrors lerobot's ``make_robot_controller``, which is not reused because it
    matches against a hardcoded list of the three controllers that shipped with
    it: a newly added one would be discovered and offered in the UI, then be
    rejected on click. Going through the discovered module path means adding a
    controller really is just dropping in a file.
    """
    modules = discover_controllers()
    if name not in modules:
        raise ValueError(f"Unknown controller: {name!r}. Available: {sorted(modules)}")
    return getattr(importlib.import_module(modules[name]), name)()