"""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.gamepad` owns the hardware, :mod:`dora_teleop.mapping` turns a
reading into a velocity, :mod:`dora_teleop.arbiter` says who is allowed to, and
:mod:`dora_teleop.ui` speaks the console's language.

**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.
"""

import argparse
import json
import time

import pyarrow as pa
from dora import Node

from dora_teleop import mapping, ui
from dora_teleop.arbiter import GAMEPAD, WEB, Arbiter
from dora_teleop.config import Config
from dora_teleop.gamepad import Gamepad, now
from dora_teleop.log import log
from dora_teleop.mapping import LAYOUT, pressed

REQUEST_ID = "request_id"
STOP_VELOCITY = [0.0] * 6


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.view = ui.ArmView()
        # The velocity the browser last asked for, and when. A jog is a stream;
        # what matters is whether the stream is still arriving.
        self.web_velocity = STOP_VELOCITY
        self.web_seen_at = 0.0
        self.clients = 0
        # 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
        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, start reading the pad, 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)
    gamepad = Gamepad(config.joystick_index)
    gamepad.start()

    try:
        for event in node:
            if event["type"] == "ERROR":
                raise RuntimeError(event["error"])
            if event["type"] != "INPUT":
                continue
            _handle(node, state, gamepad, 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_jog(node, STOP_VELOCITY)
        gamepad.shutdown()


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, gamepad: Gamepad, event: dict) -> None:
    """Route one input to whatever deals with it."""
    event_id = event["id"]
    if event_id == "tick":
        _tick(node, state, gamepad)
    elif event_id == "ui_command":
        _accept_ui_command(node, state, event)
    elif event_id == "clients":
        _accept_clients(state, event)
    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, gamepad: Gamepad) -> None:
    """Decide what the arm should be doing, and say so."""
    moment = now()
    sample = gamepad.read()

    _read_buttons(node, state, sample)

    # 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.
    if sample.active and state.arbiter.owner is None:
        state.arbiter.take(GAMEPAD, moment)
    if sample.active:
        state.arbiter.note_activity(GAMEPAD, moment)
    state.arbiter.expire_idle_gamepad(moment)

    _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.

    A velocity is only re-sent while it is non-zero: the arm node treats a
    stream of jogs as one continuous one, and the re-send exists so the
    controller's own expiry never fires mid-motion.
    """
    velocity = _wanted_velocity(state, sample, moment)

    if mapping.is_moving(velocity):
        state.jogging = True
        _send_jog(node, velocity)
        return

    if state.jogging:
        # Once, not every tick. Repeating it would keep the arm node believing a
        # jog is live and hold `robot_state` at `moving` forever.
        state.jogging = False
        _send_jog(node, STOP_VELOCITY)


def _wanted_velocity(state: Teleop, sample, moment: float) -> list[float]:
    """Return what the arm should be doing, 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.
    """
    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_VELOCITY
        return mapping.velocity(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_VELOCITY
        if moment - state.web_seen_at > state.config.deadman.jog_ttl_s:
            return STOP_VELOCITY
        return state.web_velocity

    return STOP_VELOCITY


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)

    # The stop is not gated on the lease. Somebody reaching for it is the one
    # moment when asking who owns the arm is the wrong question.
    if tapped(LAYOUT.estop):
        _estop(node, state)
    if not state.arbiter.may_drive(GAMEPAD):
        return
    if tapped(LAYOUT.grasp):
        _request(node, state, {"service": "grasp"})
    if tapped(LAYOUT.release):
        _request(node, state, {"service": "release"})
    if tapped(LAYOUT.clear):
        _request(node, state, {"service": "clear"})
    if tapped(LAYOUT.freedrive):
        _request(node, state, {"service": "freedrive", "enabled": 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 _accept_ui_command(node: Node, state: Teleop, event: dict) -> None:
    """Take one console message and do what it says."""
    raw = event["value"][0].as_py() if len(event["value"]) else None
    if raw is None:
        return
    try:
        payload = json.loads(raw) if isinstance(raw, str) else raw
    except json.JSONDecodeError:
        log("ignoring a console message that was not JSON")
        return

    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":
        # Clearing the latch is this node's business; clearing the controller's
        # error is the arm's. Both, in that order, or the console's RESET button
        # leaves one of the two still saying stop.
        state.view.estop = False
        _request(node, state, {"service": "clear"})
        return

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

    if command.type == "jog":
        state.web_velocity = command.velocity
        state.web_seen_at = moment
        state.arbiter.note_activity(WEB, moment)
    elif command.type == "jog_stop":
        state.web_velocity = STOP_VELOCITY
        state.web_seen_at = moment
    elif command.type == "home":
        _request(node, state, {"service": "home", "wait": False})
    elif command.type == "gripper":
        _request(node, state, {"service": "grasp" if command.enabled else "release"})
    elif command.type == "freedrive":
        _request(node, state, {"service": "freedrive", "enabled": 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_velocity = STOP_VELOCITY
        return
    if not state.arbiter.take(WEB, moment, preempt=command.preempt):
        _post_ui(node, ui.error(f"the {state.arbiter.owner} has the arm"))


def _accept_clients(state: Teleop, event: dict) -> None:
    """Track who is watching, because the last one leaving has to stop the arm."""
    row = event["value"][0].as_py()
    state.clients = int(row.get("connected", 0))
    if state.clients == 0:
        state.arbiter.release(WEB)
        state.web_velocity = STOP_VELOCITY


# --- 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, 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, 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_velocity = STOP_VELOCITY
    _send_jog(node, STOP_VELOCITY)
    _request(node, state, {"service": "stop"})


def _send_jog(node: Node, velocity: list[float]) -> None:
    """Send one velocity to the arm. There is no answer to a jog, ever."""
    node.send_output("jog", pa.array([{"velocity": [float(v) for v in velocity]}]))


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, ui.snapshot(state.view, state.arbiter.owner, state.jogging, time.time()))


def _post_ui(node: Node, payload: dict) -> None:
    """Hand one object to the bridge, which broadcasts it to every browser."""
    node.send_output("ui_state", pa.array([json.dumps(payload)]), metadata={"encoding": "json"})


if __name__ == "__main__":
    main()