Move majority of listbox option focus management into hooks

15 files changed, 145 insertions, 100 deletions

diff --git a/desktop/packages/mullvad-vpn/src/renderer/components/settings-listbox/components/input-option/components/input-option-trigger/InputOptionTrigger.tsx b/desktop/packages/mullvad-vpn/src/renderer/components/settings-listbox/components/input-option/components/input-option-trigger/InputOptionTrigger.tsx
index 25f71a03e0..e7343ee11a 100644
--- a/desktop/packages/mullvad-vpn/src/renderer/components/settings-listbox/components/input-option/components/input-option-trigger/InputOptionTrigger.tsx
+++ b/desktop/packages/mullvad-vpn/src/renderer/components/settings-listbox/components/input-option/components/input-option-trigger/InputOptionTrigger.tsx
@@ -46,41 +46,25 @@ export const InputOptionTrigger = ({ children, ...props }: InputOptionTriggerPro
     inputState: { value: inputValue },
   } = useInputOptionContext();
 
-  const {
-    value: selectedValue,
-    focusedValue,
-    setFocusedValue,
-    onValueChange,
-  } = useListboxContext();
+  const { value: selectedValue, onValueChange } = useListboxContext();
   const selected = value === selectedValue;
-  const focused = value === focusedValue;
-
-  const tabIndex = selected && focusedValue === undefined ? 0 : -1;
-
-  React.useEffect(() => {
-    if (focused && inputRef.current) {
-      inputRef.current.focus();
-    }
-  }, [value, focused, inputRef]);
 
   const handleClick = React.useCallback(async () => {
-    setFocusedValue(value);
     inputRef.current?.focus();
     if (!selected) {
       await onValueChange?.(inputValue);
     }
-  }, [inputRef, inputValue, onValueChange, selected, setFocusedValue, value]);
+  }, [inputRef, inputValue, onValueChange, selected]);
 
   const handleFocus = React.useCallback(() => {
-    setFocusedValue(value);
     inputRef.current?.focus();
-  }, [inputRef, setFocusedValue, value]);
+  }, [inputRef]);
 
   return (
     <StyledInputOptionTrigger
       role="option"
       aria-selected={selected}
-      tabIndex={tabIndex}
+      tabIndex={-1}
       onFocus={handleFocus}
       onClick={handleClick}
       {...props}>
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/ListboxContext.tsx b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/ListboxContext.tsx
index e7a9e89fd0..55b1a20699 100644
--- a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/ListboxContext.tsx
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/ListboxContext.tsx
@@ -4,8 +4,9 @@ import { ListboxProps } from './Listbox';
 
 type ListboxContext<T> = Pick<ListboxProps<T>, 'value' | 'onValueChange'> & {
   labelId: string;
-  focusedValue?: T;
-  setFocusedValue: React.Dispatch<React.SetStateAction<T | undefined>>;
+  optionsRef: React.RefObject<HTMLUListElement | null>;
+  focusedIndex?: number;
+  setFocusedIndex: React.Dispatch<React.SetStateAction<number | undefined>>;
 };
 
 type ListboxProviderProps<T> = Pick<ListboxContext<T>, 'value' | 'onValueChange' | 'labelId'> & {
@@ -29,15 +30,17 @@ export function ListboxProvider<T>({
   children,
 }: ListboxProviderProps<T>) {
   const TypedListboxContext = ListboxContext as React.Context<ListboxContext<T>>;
-  const [focusedValue, setFocusedValue] = React.useState<T | undefined>(undefined);
+  const [focusedIndex, setFocusedIndex] = React.useState<number | undefined>(undefined);
+  const optionsRef = React.useRef<HTMLUListElement>(null);
   return (
     <TypedListboxContext.Provider
       value={{
         value,
         onValueChange,
         labelId,
-        focusedValue,
-        setFocusedValue,
+        focusedIndex,
+        setFocusedIndex,
+        optionsRef,
       }}>
       {children}
     </TypedListboxContext.Provider>
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-option/components/listbox-option-trigger/ListboxOptionTrigger.tsx b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-option/components/listbox-option-trigger/ListboxOptionTrigger.tsx
index 9a587e826d..383a0aac32 100644
--- a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-option/components/listbox-option-trigger/ListboxOptionTrigger.tsx
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-option/components/listbox-option-trigger/ListboxOptionTrigger.tsx
@@ -64,16 +64,8 @@ export const ListboxOptionTrigger = ({ children, ...props }: ListboxOptionTrigge
   const { disabled } = useListItemContext();
   const triggerRef = React.useRef<HTMLLIElement>(null);
 
-  const {
-    value: selectedValue,
-    onValueChange,
-    focusedValue,
-    setFocusedValue,
-  } = useListboxContext();
+  const { value: selectedValue, onValueChange } = useListboxContext();
   const selected = value === selectedValue;
-  const focused = value === focusedValue;
-
-  const tabIndex = !disabled && selected && focusedValue === undefined ? 0 : -1;
 
   const handleTriggerClick = React.useCallback(async () => {
     if (onValueChange) {
@@ -81,17 +73,6 @@ export const ListboxOptionTrigger = ({ children, ...props }: ListboxOptionTrigge
     }
   }, [onValueChange, value]);
 
-  React.useEffect(() => {
-    if (focused && triggerRef.current) {
-      triggerRef.current.focus();
-    }
-  }, [value, focused]);
-
-  const onFocus = React.useCallback(() => {
-    if (focused) return;
-    setFocusedValue(value);
-  }, [focused, setFocusedValue, value]);
-
   const handleClick = !disabled ? handleTriggerClick : undefined;
 
   const handleKeyDown = React.useCallback(
@@ -113,10 +94,9 @@ export const ListboxOptionTrigger = ({ children, ...props }: ListboxOptionTrigge
       role="option"
       aria-selected={selected}
       aria-disabled={disabled}
-      tabIndex={tabIndex}
+      tabIndex={-1}
       onClick={handleClick}
       onKeyDown={handleKeyDown}
-      onFocus={onFocus}
       $disabled={disabled}
       {...props}>
       {children}
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/ListboxOptions.tsx b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/ListboxOptions.tsx
index 3ceccfad54..6af21a10cb 100644
--- a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/ListboxOptions.tsx
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/ListboxOptions.tsx
@@ -1,18 +1,34 @@
 import React from 'react';
 
 import { useListboxContext } from '../../';
-import { useChildrenValues, useHandleKeyboardNavigation } from './hooks';
+import { useHandleKeyboardNavigation } from './hooks';
+import { getInitialOption, getOptions } from './utils';
 
 export type ListboxOptionsProps = {
   children: React.ReactNode[];
 };
 
 export function ListboxOptions({ children }: ListboxOptionsProps) {
-  const { labelId, setFocusedValue } = useListboxContext();
-  const ref = React.useRef<HTMLUListElement>(null);
-  const values = useChildrenValues(children);
+  const { labelId, optionsRef, setFocusedIndex } = useListboxContext();
+  const [tabIndex, setTabIndex] = React.useState<number>(0);
+
+  const handleFocus = React.useCallback(
+    (event: React.FocusEvent) => {
+      if (!optionsRef.current?.isSameNode(event.target)) return;
+
+      const options = getOptions(optionsRef.current);
+
+      const initialOption = getInitialOption(options);
+      if (initialOption) {
+        setTabIndex(-1);
+        initialOption.focus();
+      }
+    },
+    [optionsRef],
+  );
+
+  const handleKeyboardNavigation = useHandleKeyboardNavigation();
 
-  const handleKeyboardNavigation = useHandleKeyboardNavigation(values);
   const onKeyDown = React.useCallback(
     (event: React.KeyboardEvent) => {
       handleKeyboardNavigation(event);
@@ -20,21 +36,29 @@ export function ListboxOptions({ children }: ListboxOptionsProps) {
     [handleKeyboardNavigation],
   );
 
-  const onBlur = React.useCallback(
-    (e: React.FocusEvent<HTMLUListElement>) => {
-      const container = ref.current;
-      const nextFocus = e.relatedTarget as Node | null;
+  const handleBlur = React.useCallback(
+    (event: React.FocusEvent<HTMLUListElement>) => {
+      const container = optionsRef.current;
+      const nextFocus = event.relatedTarget as Node | null;
 
       // If focus moves outside the listbox
       if (!container || !nextFocus || !container.contains(nextFocus)) {
-        setFocusedValue(undefined);
+        setFocusedIndex(undefined);
+        setTabIndex(0);
       }
     },
-    [setFocusedValue],
+    [optionsRef, setFocusedIndex],
   );
 
   return (
-    <ul ref={ref} role="listbox" onKeyDown={onKeyDown} onBlur={onBlur} aria-labelledby={labelId}>
+    <ul
+      ref={optionsRef}
+      role="listbox"
+      aria-labelledby={labelId}
+      onKeyDown={onKeyDown}
+      onBlur={handleBlur}
+      onFocus={handleFocus}
+      tabIndex={tabIndex}>
       {children}
     </ul>
   );
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/index.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/index.ts
index c112a433de..64b446899b 100644
--- a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/index.ts
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/index.ts
@@ -1,2 +1 @@
 export * from './useHandleKeyboardNavigation';
-export * from './useChildrenValues';
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/useChildrenValues.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/useChildrenValues.ts
deleted file mode 100644
index 29e58da9cd..0000000000
--- a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/useChildrenValues.ts
+++ /dev/null
@@ -1,17 +0,0 @@
-import React from 'react';
-
-export function useChildrenValues(children: React.ReactNode[]): string[] {
-  return React.useMemo(() => {
-    const values: string[] = [];
-
-    React.Children.forEach(children, (child) => {
-      if (React.isValidElement<{ value?: string }>(child)) {
-        if ('value' in child.props && typeof child.props.value !== 'undefined') {
-          values.push(child.props.value);
-        }
-      }
-    });
-
-    return values;
-  }, [children]);
-}
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/useFocusOptionByIndex.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/useFocusOptionByIndex.ts
new file mode 100644
index 0000000000..8e75d70912
--- /dev/null
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/useFocusOptionByIndex.ts
@@ -0,0 +1,17 @@
+import React from 'react';
+
+import { useListboxContext } from '../../../ListboxContext';
+import { getOptions } from '../utils';
+
+export const useFocusOptionByIndex = () => {
+  const { setFocusedIndex, optionsRef } = useListboxContext();
+  return React.useCallback(
+    (index: number) => {
+      const options = getOptions(optionsRef.current);
+      setFocusedIndex(index);
+      const option = options[index];
+      option.focus();
+    },
+    [optionsRef, setFocusedIndex],
+  );
+};
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/useGetInitialFocusIndex.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/useGetInitialFocusIndex.ts
new file mode 100644
index 0000000000..43205589cd
--- /dev/null
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/useGetInitialFocusIndex.ts
@@ -0,0 +1,19 @@
+import React from 'react';
+
+import { useListboxContext } from '../../../ListboxContext';
+import { getOptions, getSelectedOptionIndex } from '../utils';
+
+export const useGetInitialFocusIndex = () => {
+  const { focusedIndex, optionsRef } = useListboxContext();
+  return React.useCallback(() => {
+    const options = getOptions(optionsRef.current);
+    if (focusedIndex !== undefined) {
+      return focusedIndex;
+    }
+    const selectedOptionIndex = getSelectedOptionIndex(options);
+    if (selectedOptionIndex !== -1) {
+      return selectedOptionIndex;
+    }
+    return 0;
+  }, [focusedIndex, optionsRef]);
+};
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/useHandleKeyboardNavigation.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/useHandleKeyboardNavigation.ts
index 335b2d12f0..ba7b61bd67 100644
--- a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/useHandleKeyboardNavigation.ts
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/hooks/useHandleKeyboardNavigation.ts
@@ -1,43 +1,41 @@
 import React from 'react';
 
 import { useListboxContext } from '../../../';
+import { getOptions } from '../utils';
+import { useFocusOptionByIndex } from './useFocusOptionByIndex';
+import { useGetInitialFocusIndex } from './useGetInitialFocusIndex';
 
-export const useHandleKeyboardNavigation = <T>(options: T[]) => {
-  const { value: selectedValue, focusedValue, setFocusedValue } = useListboxContext<T>();
+export const useHandleKeyboardNavigation = () => {
+  const { optionsRef } = useListboxContext();
+  const getInitialFocusIndex = useGetInitialFocusIndex();
+  const focusOptionByIndex = useFocusOptionByIndex();
 
   return React.useCallback(
     (event: React.KeyboardEvent) => {
-      let value: T;
-      if (focusedValue !== undefined) value = focusedValue;
-      else if (selectedValue !== undefined) value = selectedValue;
-      else value = options[0];
+      const options = getOptions(optionsRef.current);
+
+      const initialFocusedIndex = getInitialFocusIndex();
 
-      // Use roving tabindex to determine the next focusable option
-      let nextValue: T | undefined;
       if (event.key === 'ArrowUp') {
         event.preventDefault();
-        const focusedOptionIndex = options.findIndex((option) => option === value);
-        if (focusedOptionIndex <= 0) return;
-        const nextOptionIndex = focusedOptionIndex - 1;
-        nextValue = options[nextOptionIndex];
+        if (initialFocusedIndex > 0) {
+          const newFocusedIndex = initialFocusedIndex - 1;
+          focusOptionByIndex(newFocusedIndex);
+        }
       } else if (event.key === 'ArrowDown') {
         event.preventDefault();
-        const focusedOptionIndex = options.findIndex((option) => option === value);
-        if (focusedOptionIndex >= options.length - 1) return;
-        const nextOptionIndex = focusedOptionIndex + 1;
-        if (nextOptionIndex === -1) return;
-        nextValue = options[nextOptionIndex];
+        if (initialFocusedIndex < options.length - 1) {
+          const newFocusedIndex = initialFocusedIndex + 1;
+          focusOptionByIndex(newFocusedIndex);
+        }
       } else if (event.key === 'Home') {
         event.preventDefault();
-        nextValue = options[0];
+        focusOptionByIndex(0);
       } else if (event.key === 'End') {
         event.preventDefault();
-        nextValue = options[options.length - 1];
-      }
-      if (nextValue !== undefined) {
-        setFocusedValue(nextValue);
+        focusOptionByIndex(options.length - 1);
       }
     },
-    [focusedValue, selectedValue, options, setFocusedValue],
+    [focusOptionByIndex, getInitialFocusIndex, optionsRef],
   );
 };
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-initial-option.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-initial-option.ts
new file mode 100644
index 0000000000..e2278b5237
--- /dev/null
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-initial-option.ts
@@ -0,0 +1,10 @@
+import { getSelectedOption } from './get-selected-option';
+
+export const getInitialOption = (options: HTMLElement[]) => {
+  const selectedOption = getSelectedOption(options);
+  if (selectedOption) {
+    return selectedOption;
+  }
+
+  return options.length ? options[0] : undefined;
+};
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-is-option-selected.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-is-option-selected.ts
new file mode 100644
index 0000000000..6433e20232
--- /dev/null
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-is-option-selected.ts
@@ -0,0 +1,3 @@
+export const getIsOptionSelected = (option: HTMLElement) => {
+  return option.getAttribute('aria-selected') === 'true';
+};
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-options.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-options.ts
new file mode 100644
index 0000000000..d37647936c
--- /dev/null
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-options.ts
@@ -0,0 +1,11 @@
+export const getOptions = (container: HTMLElement | null) => {
+  const options = container?.querySelectorAll<HTMLElement>(
+    '[role="option"]:not([aria-disabled="true"])',
+  );
+
+  if (options) {
+    return Array.from(options);
+  }
+
+  return [];
+};
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-selected-option-index.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-selected-option-index.ts
new file mode 100644
index 0000000000..2534ef46d1
--- /dev/null
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-selected-option-index.ts
@@ -0,0 +1,5 @@
+import { getIsOptionSelected } from './get-is-option-selected';
+
+export const getSelectedOptionIndex = (options: HTMLElement[]) => {
+  return options.findIndex((option) => getIsOptionSelected(option));
+};
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-selected-option.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-selected-option.ts
new file mode 100644
index 0000000000..dfed49aae7
--- /dev/null
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/get-selected-option.ts
@@ -0,0 +1,5 @@
+import { getIsOptionSelected } from './get-is-option-selected';
+
+export const getSelectedOption = (options: HTMLElement[]) => {
+  return options.find((option) => getIsOptionSelected(option));
+};
diff --git a/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/index.ts b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/index.ts
new file mode 100644
index 0000000000..cbd0929620
--- /dev/null
+++ b/desktop/packages/mullvad-vpn/src/renderer/lib/components/listbox/components/listbox-options/utils/index.ts
@@ -0,0 +1,4 @@
+export * from './get-initial-option';
+export * from './get-options';
+export * from './get-selected-option-index';
+export * from './get-selected-option';