kvm/ui/src/hooks/useKeyboard.ts

import { useCallback } from "react";

import { KeysDownState, HidState, useHidStore, RTCState, useRTCStore, hidKeyBufferSize, hidErrorRollOver } from "@/hooks/stores";
import { JsonRpcResponse, useJsonRpc } from "@/hooks/useJsonRpc";
import { hidKeyToModifierMask, keys, modifiers } from "@/keyboardMappings";

export default function useKeyboard() {
  const [send] = useJsonRpc();

  const rpcDataChannel = useRTCStore((state: RTCState) => state.rpcDataChannel);

  const keysDownState = useHidStore((state: HidState) => state.keysDownState);
  const setKeysDownState = useHidStore((state: HidState) => state.setKeysDownState);

  const keyPressAvailable = useHidStore((state: HidState) => state.keyPressAvailable);
  const setKeyPressAvailable = useHidStore((state: HidState) => state.setKeyPressAvailable);

  const sendKeyboardEvent = useCallback(
    (state: KeysDownState) => {
      if (rpcDataChannel?.readyState !== "open") return;

      console.debug(`Send keyboardReport keys: ${state.keys}, modifier: ${state.modifier}`);
      send("keyboardReport", { keys: state.keys, modifier: state.modifier }, (resp: JsonRpcResponse) => {
        if ("error" in resp) {
          console.error(`Failed to send keyboard report ${state}`, resp.error);
        } else {
          const keysDownState = resp.result as KeysDownState;

          if (keysDownState) {
            // new devices return the keyDownState, so we can use it to update the state
            setKeysDownState(keysDownState);
            setKeyPressAvailable(true); // if they returned a keysDownState, we know they also support keyPressReport
          } else {
            // old devices do not return the keyDownState, so we just pretend they accepted what we sent
            setKeysDownState(state);
            // and we shouldn't set keyPressAvailable here because we don't know if they support it
          }
        }
      });
    },
    [rpcDataChannel?.readyState, send, setKeyPressAvailable, setKeysDownState],
  );

  const sendKeypressEvent = useCallback(
    (key: number, press: boolean) => {
      if (rpcDataChannel?.readyState !== "open") return;

      console.debug(`Send keypressEvent key: ${key}, press: ${press}`);
      send("keypressReport", { key, press }, (resp: JsonRpcResponse) => {
        if ("error" in resp) {
          // -32601 means the method is not supported
          if (resp.error.code === -32601) {
            // if we don't support key press report, we need to disable all that handling
            console.error("Failed calling keypressReport, switching to local handling", resp.error);
            setKeyPressAvailable(false);
          } else {
            console.error(`Failed to send key ${key} press: ${press}`, resp.error);
          }
        } else {
          const keysDownState = resp.result as KeysDownState;

          if (keysDownState) {
            setKeysDownState(keysDownState);
            // we don't need to set keyPressAvailable here, because it's already true or we never landed here
          }
        }
      });
    },
    [rpcDataChannel?.readyState, send, setKeyPressAvailable, setKeysDownState],
  );

  const resetKeyboardState = useCallback(() => {
    console.debug("Resetting keyboard state");
    keysDownState.keys.fill(0); // Reset the keys buffer to zeros
    keysDownState.modifier = 0; // Reset the modifier state to zero
    sendKeyboardEvent(keysDownState);
  }, [keysDownState, sendKeyboardEvent]);

  const executeMacro = async (steps: { keys: string[] | null; modifiers: string[] | null; delay: number }[]) => {
    for (const [index, step] of steps.entries()) {
      const keyValues = (step.keys || []).map(key => keys[key]).filter(Boolean);
      const modifierMask: number = (step.modifiers || []).map(mod => modifiers[mod]).reduce((acc, val) => acc + val, 0);

      // If the step has keys and/or modifiers, press them and hold for the delay
      if (keyValues.length > 0 || modifierMask > 0) {
        sendKeyboardEvent({ keys: keyValues, modifier: modifierMask });
        await new Promise(resolve => setTimeout(resolve, step.delay || 50));

        resetKeyboardState();
      } else {
        // This is a delay-only step, just wait for the delay amount
        await new Promise(resolve => setTimeout(resolve, step.delay || 50));
      }

      // Add a small pause between steps if not the last step
      if (index < steps.length - 1) {
        await new Promise(resolve => setTimeout(resolve, 10));
      }
    }
  };

  // this code exists because we have devices that don't support the keysPress api yet (not current)
  // so we mirror the device-side code here to keep track of the keyboard state
  function handleKeyLocally(state: KeysDownState, key: number, press: boolean): KeysDownState {
    const keys = state.keys;
    let modifiers = state.modifier;
    const modifierMask = hidKeyToModifierMask[key] || 0;

    if (modifierMask !== 0) {
      console.debug(`Handling modifier key: ${key}, press: ${press}, current modifiers: ${modifiers}, modifier mask: ${modifierMask}`);
      if (press) {
        modifiers |= modifierMask;
      } else {
        modifiers &= ~modifierMask;
      }
    } else {
      // handle other keys that are not modifier keys by placing or removing them
      // from the key buffer since the buffer tracks currently pressed keys
      let overrun = true;
      for (let i = 0; i < hidKeyBufferSize && overrun; i++) {
        // If we find the key in the buffer the buffer, we either remove it (if press is false)
        // or do nothing (if down is true) because the buffer tracks currently pressed keys
        // and if we find a zero byte, we can place the key there (if press is true)
        if (keys[i] == key || keys[i] == 0) {
          if (press) {
            keys[i] = key // overwrites the zero byte or the same key if already pressed
          } else {
            // we are releasing the key, remove it from the buffer
            if (keys[i] != 0) {
              keys.splice(i, 1);
              keys.push(0); // add a zero at the end
            }
          }
          overrun = false // We found a slot for the key
        }

        // If we reach here it means we didn't find an empty slot or the key in the buffer
        if (overrun) {
          if (press) {
            console.warn(`keyboard buffer overflow, key: ${key} not added`);
            // Fill all key slots with ErrorRollOver (0x01) to indicate overflow
            keys.length = 6;
            keys.fill(hidErrorRollOver);
          } else {
            // If we are releasing a key, and we didn't find it in a slot, who cares?
            console.debug(`key ${key} not found in buffer, nothing to release`)
          }
        }
      }
    }
    return { modifier: modifiers, keys };
  }

  const handleKeyPress = useCallback(
    (key: number, press: boolean) => {
      if (rpcDataChannel?.readyState !== "open") return;
  
      if (keyPressAvailable) {
        // if the keyPress api is available, we can just send the key press event
        sendKeypressEvent(key, press);
        // if keyPress api is STILL available, we don't need to handle the key locally
        if (keyPressAvailable) return;
      }

      // if the keyPress api is not available, we need to handle the key locally
      const downState = handleKeyLocally(keysDownState, key, press);
      setKeysDownState(downState);

      // then we send the full state
      sendKeyboardEvent(downState);
    },
    [keyPressAvailable, keysDownState, rpcDataChannel?.readyState, sendKeyboardEvent, sendKeypressEvent, setKeysDownState],
  );

  return { handleKeyPress, resetKeyboardState, executeMacro };
}