"""The dora event loop: routing, the deadman, and nothing that can block.

The loop owns the one decision this node exists to make — whether a velocity
reaches the arm this tick, and from whom. Everything else is delegated:
:mod:`dora_teleop.joy` takes the pad off the wire, :mod:`dora_teleop.mapping`
turns a reading into a velocity, :mod:`dora_teleop.arbiter` says who is allowed
to, :mod:`dora_teleop.ui` speaks the console's language and
:mod:`dora_teleop.server` serves it. **No hardware is touched here at all** —
the pad is `dora-joy` and the arm is `dora-ulite6`.

**The safety-critical part is `_drive`.** An arm moving under a velocity command
keeps moving until told otherwise, so the question is never "should it start"
but "has anything happened that means it must stop". Three things mean that: the
source went quiet, the source went away, or this process is going away. All
three end in the same place — one all-zero velocity, sent once.

What goes out is the *change*: a velocity is a set-point the arm holds, not a
frame in a sequence, so an unchanged one is repeated only as often as the arm's
own expiry demands. A change — a new direction, or the stop — is never made to
wait for that timer.
"""

import argparse
import json
import time

import pyarrow as pa
from dora import Node

from dora_teleop import joy, mapping, message, server, ui
from dora_teleop.arbiter import GAMEPAD, WEB, Arbiter
from dora_teleop.config import Config
from dora_teleop.joy import JoyInput, now
from dora_teleop.log import log
from dora_teleop.mapping import LAYOUT, pressed

REQUEST_ID = "request_id"
# Still, in the normalised units both sources speak.
STOP_AXES = [0.0] * 6

# The verbs this node publishes on `command`. Robot-agnostic on purpose: a
# consumer that grows a two-finger gripper where the last one had a vacuum
# should need no change here, and neither should this node.
GRIPPER_CLOSE = "gripper_close"
GRIPPER_OPEN = "gripper_open"
STOP = "stop"
HOME = "home"
CLEAR = "clear"
FREEDRIVE = "freedrive"

# The same verbs as `dora-ulite6` names them, for the deprecated `request` port.
# This table is the whole of what used to be spread through the loop, and it is
# what a robot speaking the standard lets us delete.
LEGACY_REQUEST = {
    GRIPPER_CLOSE: {"service": "grasp"},
    GRIPPER_OPEN: {"service": "release"},
    STOP: {"service": "stop"},
    CLEAR: {"service": "clear"},
    HOME: {"service": "home", "wait": False},
}


class Teleop:
    """Everything the loop remembers between ticks."""

    def __init__(self, config: Config):
        self.config = config
        self.arbiter = Arbiter(config.deadman.gamepad_idle_release_s)
        self.joy = JoyInput(config.deadman.joy_ttl_s)
        self.view = ui.ArmView()
        # What the browser last asked for, and when. A jog is a stream; what
        # matters is whether the stream is still arriving. Normalised, like the
        # pad's — the console does not know the robot's units.
        self.web_axes = STOP_AXES
        self.web_seen_at = 0.0
        # How many browsers are on the console right now, read from the server
        # every tick. The last one leaving has to stop the arm, and a socket
        # that went away sends nothing to say so.
        self.clients = 0
        self.console: server.Console | None = None
        # Whether the arm is currently being driven by us. Used to send the stop
        # once rather than every tick: a stream of zeroes would keep the arm
        # node believing a jog is live and hold `robot_state` at `moving`.
        self.jogging = False
        # The last velocity actually put on the wire, and when. `_drive` sends
        # only when what it wants differs from this, so the arm is told about
        # changes rather than told the same thing 20 times a second.
        self.sent_axes = STOP_AXES
        self.sent_at = 0.0
        self.buttons: list[bool] = []
        self.requests = 0
        self.polled_at = 0.0

    def next_request_id(self, service: str) -> str:
        """Mint a correlation id we will recognise coming back."""
        self.requests += 1
        return f"teleop-{service}-{self.requests}"


def main() -> None:
    """Connect to the dataflow and run until dora stops us."""
    args = _parse_args()
    config = Config.from_env()

    # Passing the name lets the node also run standalone as a dynamic node; when
    # dora spawns it, the id from the dataflow wins.
    node = Node(args.name)
    state = Teleop(config)
    if config.console.enabled:
        state.console = server.Console(config.console.host, config.console.port)
        state.console.start()

    try:
        for event in node:
            if event["type"] == "ERROR":
                raise RuntimeError(event["error"])
            if event["type"] != "INPUT":
                continue
            _handle(node, state, event)
    finally:
        # A motion lives in the arm's controller, not in this process. Leaving
        # without zeroing the velocity leaves it travelling until the
        # controller's own expiry catches up, which is the one moment nobody is
        # watching for it.
        _send_motion(node, state, STOP_AXES)
        if state.console is not None:
            state.console.stop()


def _parse_args() -> argparse.Namespace:
    """Read the node name, for when dora did not spawn us."""
    parser = argparse.ArgumentParser(description="Drive a robot arm from a gamepad or a browser.")
    parser.add_argument(
        "--name",
        default="teleop",
        help="The name of the node in the dataflow, used when running as a dynamic node.",
    )
    return parser.parse_args()


def _handle(node: Node, state: Teleop, event: dict) -> None:
    """Route one input to whatever deals with it."""
    event_id = event["id"]
    if event_id == "tick":
        _tick(node, state)
    elif event_id == "joy":
        state.joy.accept(joy.decode(event), now())
    elif event_id == "robot_state":
        state.view.state = event["value"][0].as_py()
    elif event_id == "tcp_pose":
        state.view.pose = event["value"].to_pylist()
    elif event_id == "joint_pose":
        state.view.joints = event["value"].to_pylist()
    elif event_id == "response":
        _accept_response(node, state, event)
    elif event_id == "motion_state":
        _accept_motion_state(node, state, event)
    else:
        log(f"ignoring unexpected input `{event_id}`")


# --- the tick, where the arm is actually driven ---


def _tick(node: Node, state: Teleop) -> None:
    """Decide what the arm should be doing, and say so."""
    moment = now()
    # Before anything else: what the browser asked for since the last tick is
    # part of this tick's picture, and the client count it implies is what the
    # web deadman reads below.
    _drain_console(node, state)
    # Absent covers three things now: no pad plugged in, a pad unplugged
    # mid-jog, and `dora-joy` no longer sending. The arm cannot tell them apart
    # and does not need to.
    sample = state.joy.read(moment)

    # Who owns the arm is settled before anything they pressed is acted on.
    # The other order drops the first press after the pad has been idle: the
    # lease is not held yet when the buttons are read, and by the next tick the
    # button is held rather than newly down, so the edge is gone for good.
    #
    # A pad taking the arm the moment a stick moves is what keeps an operator
    # from having to reach for a browser to use the thing in their hands.
    active = mapping.is_active(sample)
    if active and state.arbiter.owner is None:
        state.arbiter.take(GAMEPAD, moment)
    if active:
        state.arbiter.note_activity(GAMEPAD, moment)
    state.arbiter.expire_idle_gamepad(moment)

    _read_buttons(node, state, sample)

    _drive(node, state, sample, moment)
    _poll_state(node, state, moment)
    _publish(node, state)


def _drive(node: Node, state: Teleop, sample, moment: float) -> None:
    """Send the velocity the owner is asking for, or the stop that ends it.

    Changes go out at once; an unchanged velocity is repeated only on
    `jog_keepalive_s`, which exists solely so the arm's own velocity expiry
    never fires mid-motion. Repeating it every tick instead made the controller
    re-plan the same motion twenty times a second, which an operator feels as
    the arm moving in steps rather than travelling.

    The stop is a change like any other, so releasing a button is acted on the
    tick it happens and never waits for the repeat.
    """
    axes = _wanted_axes(state, sample, moment)

    if axes != state.sent_axes:
        # Includes the stop: going to all-zero is a change, sent once, and the
        # `jogging` flag then keeps it from being repeated. A stream of zeroes
        # would hold the arm node's `robot_state` at `moving` forever.
        state.jogging = mapping.is_moving(axes)
        _send_motion(node, state, axes)
        return

    if state.jogging and moment - state.sent_at >= state.config.deadman.jog_keepalive_s:
        _send_motion(node, state, axes)


def _wanted_axes(state: Teleop, sample, moment: float) -> list[float]:
    """Return what the robot should be doing, normalised, which is usually nothing.

    Every path that is not an owner actively asking for motion returns a stop.
    That is deliberate: the default has to be "still", so that a case nobody
    thought of fails by halting rather than by carrying on.
    """
    # First, and above the lease. A stop somebody had to reach for is not undone
    # by the next thumb movement: it holds until it is deliberately cleared.
    if state.view.estop:
        return STOP_AXES

    owner = state.arbiter.owner
    if owner == GAMEPAD:
        # An unplugged pad is not a centred pad. Stopping here rather than
        # waiting out the TTL is the difference between a dropped controller
        # costing 10 ms of travel and costing 400.
        if not sample.present:
            return STOP_AXES
        return mapping.from_pad(sample.axes, sample.buttons, state.config)

    if owner == WEB:
        # The console refreshes its jog every 150 ms; the browser going quiet,
        # or the tab closing, has to stop the arm without anything being sent.
        if state.clients == 0:
            return STOP_AXES
        if moment - state.web_seen_at > state.config.deadman.jog_ttl_s:
            return STOP_AXES
        return state.web_axes

    return STOP_AXES


def _read_buttons(node: Node, state: Teleop, sample) -> None:
    """Fire the services a gamepad button asks for, on the press only.

    Edge-triggered, because these are one-shot requests rather than a held
    velocity: holding `grasp` should not send a request every 50 ms.
    """
    previous, state.buttons = state.buttons, list(sample.buttons)
    if not sample.present:
        return

    def tapped(index: int) -> bool:
        return pressed(sample.buttons, index) and not pressed(previous, index)

    # Neither the stop nor its release is gated on the lease. Somebody reaching
    # for the stop is the one moment when asking who owns the arm is the wrong
    # question, and a latch that cannot be released without first taking the arm
    # back is a trap — an idle pad hands the lease away two seconds after the
    # stop it just made. The console's RESET is ungated for the same reason.
    if tapped(LAYOUT.estop):
        _estop(node, state)
    if tapped(LAYOUT.clear):
        _clear(node, state)
    if not state.arbiter.may_drive(GAMEPAD):
        return
    if tapped(LAYOUT.grasp):
        _act(node, state, GRIPPER_CLOSE)
    if tapped(LAYOUT.release):
        _act(node, state, GRIPPER_OPEN)
    if tapped(LAYOUT.freedrive):
        _freedrive(node, state, not state.view.freedrive_on)


def _poll_state(node: Node, state: Teleop, moment: float) -> None:
    """Ask the arm for what telemetry does not carry.

    Whether the tool is engaged and whether freedrive is on only come back from
    `get_state`, which the arm node answers off its report cache and so can
    answer during a motion.
    """
    if moment - state.polled_at < state.config.deadman.state_poll_s:
        return
    state.polled_at = moment
    _request(node, state, {"service": "get_state"})


# --- what the console asks for ---


def _drain_console(node: Node, state: Teleop) -> None:
    """Act on everything the browsers sent, then note how many are still there."""
    if state.console is None:
        return
    for payload in state.console.take():
        _accept_ui_command(node, state, payload)
    _set_clients(state, state.console.clients)


def _accept_ui_command(node: Node, state: Teleop, payload: dict) -> None:
    """Take one console message and do what it says."""
    command = ui.command_of(payload)
    if command is None:
        return
    moment = now()

    if command.type == "estop":
        _estop(node, state)
        return
    if command.type == "control":
        _control(node, state, command, moment)
        return
    if command.type == "reset":
        _clear(node, state)
        return

    if not state.arbiter.may_drive(WEB):
        _post_ui(node, state, ui.error("take control first"))
        return

    if command.type == "jog":
        state.web_axes = mapping.from_console(command.axes)
        state.web_seen_at = moment
        state.arbiter.note_activity(WEB, moment)
    elif command.type == "jog_stop":
        state.web_axes = STOP_AXES
        state.web_seen_at = moment
    elif command.type == "home":
        _act(node, state, HOME)
    elif command.type == "gripper":
        _act(node, state, GRIPPER_CLOSE if command.enabled else GRIPPER_OPEN)
    elif command.type == "freedrive":
        _freedrive(node, state, command.enabled)
    elif command.type == "move":
        _request(node, state, {"service": "move_to", "wait": False, **command.row})


def _control(node: Node, state: Teleop, command, moment: float) -> None:
    """Hand the arm to the console, or take it back."""
    if command.action == "release":
        state.arbiter.release(WEB)
        state.web_axes = STOP_AXES
        return
    if not state.arbiter.take(WEB, moment, preempt=command.preempt):
        _post_ui(node, state, ui.error(f"the {state.arbiter.owner} has the arm"))


def _set_clients(state: Teleop, count: int) -> None:
    """Track who is watching, because the last one leaving has to stop the arm."""
    state.clients = count
    if count == 0:
        state.arbiter.release(WEB)
        state.web_axes = STOP_AXES


# --- what the arm says back ---


def _accept_response(node: Node, state: Teleop, event: dict) -> None:
    """Absorb one answer, and tell the console about it if it was bad news."""
    row = event["value"][0].as_py()
    if row.get("service") == "get_state":
        state.view.take_get_state(row)
        return
    if not row.get("ok", True):
        message = row.get("message", "")
        log(f"{row.get('service')} failed: {message}")
        _post_ui(node, state, ui.error(message))


def _accept_motion_state(node: Node, state: Teleop, event: dict) -> None:
    """Report a motion that failed, which a non-blocking request learns no other way."""
    row = event["value"][0].as_py()
    if row.get("state") == "failed":
        _post_ui(node, state, ui.error(row.get("message", "the motion failed")))


# --- talking to the dataflow ---


def _estop(node: Node, state: Teleop) -> None:
    """Halt the arm, from any source, held lease or not."""
    log("e-stop")
    state.view.estop = True
    state.jogging = False
    state.web_axes = STOP_AXES
    _send_motion(node, state, STOP_AXES)
    _act(node, state, STOP)


def _clear(node: Node, state: Teleop) -> None:
    """Release the latch and clear the controller's error, in that order.

    Both, or one of the two is still saying stop: the latch is this node's
    business and the error is the arm's. The console's RESET and the pad's
    `clear` button are the same act and go through here, so a pad-only
    dataflow is not left with a latch nothing can release.
    """
    log("cleared")
    state.view.estop = False
    _act(node, state, CLEAR)


def _send_motion(node: Node, state: Teleop, axes: list[float]) -> None:
    """Send one motion command, on the standard port and the deprecated one.

    There is no answer to either, ever — a jog is a stream that expires on the
    consumer's own TTL, so the sender keeps sending and nobody waits.

    `twist` is the standard: SI, framed, what a robot node should take. `jog`
    is `dora-ulite6`'s mm/s and deg/s, published alongside until that node
    learns the standard. Both carry the same intent, so wiring either is safe;
    wiring both would drive the arm twice.

    This always sends. The decision not to is `_drive`'s, so that the e-stop and
    the shutdown stop — neither of which goes through it — can never be
    suppressed by having sent the same thing already.
    """
    config = state.config
    node.send_output("twist", message.pack_twist(mapping.twist(axes, config)))
    node.send_output("jog", message.pack_jog(mapping.scale(axes, config)))
    state.sent_axes = list(axes)
    state.sent_at = now()


def _act(node: Node, state: Teleop, action: str, value: float | None = None) -> None:
    """Ask for one discrete thing, on the standard port and the deprecated one.

    The verb is robot-agnostic; `LEGACY_REQUEST` is the only place that knows
    what `dora-ulite6` calls it. A robot speaking the standard needs no entry
    there, which is how that table eventually goes away.
    """
    node.send_output("command", message.pack_command(action, value))
    row = LEGACY_REQUEST.get(action)
    if row is not None:
        _request(node, state, dict(row))


def _freedrive(node: Node, state: Teleop, enabled: bool) -> None:
    """Toggle freedrive. Its argument rides as the verb's value."""
    _act(node, state, FREEDRIVE, value=float(enabled))
    _request(node, state, {"service": "freedrive", "enabled": enabled})


def _request(node: Node, state: Teleop, row: dict) -> None:
    """Call one service on the arm under an id we will recognise coming back."""
    node.send_output("request", pa.array([row]), metadata={REQUEST_ID: state.next_request_id(row["service"])})


def _publish(node: Node, state: Teleop) -> None:
    """Send the console its picture of the arm."""
    _post_ui(node, state, ui.snapshot(state.view, state.arbiter.owner, state.jogging, time.time()))


def _post_ui(node: Node, state: Teleop, payload: dict) -> None:
    """Send one object to every browser, and out on the wire for anyone watching.

    `ui_state` carries no browser any more — the sockets are ours. It stays
    published because it is the one port that says what an operator is being
    shown, which is what `dora topic echo` and the bench read it for.
    """
    if state.console is not None:
        state.console.broadcast(payload)
    node.send_output("ui_state", pa.array([json.dumps(payload)]), metadata={"encoding": "json"})


if __name__ == "__main__":
    main()