From 3bc8d30190fa818e16c4c8c6a2e57cbf507470b2 Mon Sep 17 00:00:00 2001
From: Mykhailo Kravchenko <rezkiy37@gmail.com>
Date: Fri, 15 Sep 2023 14:00:02 +0200
Subject: [PATCH] fix focus on different scenarios

---
 src/components/CategoryPicker/index.js        | 26 ++++++++++---------
 .../OptionsSelector/BaseOptionsSelector.js    |  4 +--
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/src/components/CategoryPicker/index.js b/src/components/CategoryPicker/index.js
index 91c7e82e7887..4ee429a2e389 100644
--- a/src/components/CategoryPicker/index.js
+++ b/src/components/CategoryPicker/index.js
@@ -34,21 +34,23 @@ function CategoryPicker({policyCategories, reportID, iouType, iou, policyRecentl
         ];
     }, [iou.category]);
 
+    const sections = useMemo(() => {
+        const {categoryOptions} = OptionsListUtils.getNewChatOptions({}, {}, [], searchValue, selectedOptions, [], false, false, true, policyCategories, policyRecentlyUsedCategories, false);
+
+        return categoryOptions;
+    }, [policyCategories, policyRecentlyUsedCategories, searchValue, selectedOptions]);
+
     const initialFocusedIndex = useMemo(() => {
-        if (isCategoriesCountBelowThreshold && selectedOptions.length > 0) {
-            return _.chain(policyCategories)
-                .values()
-                .findIndex((category) => category.name === selectedOptions[0].name, true)
-                .value();
-        }
+        let categoryInitialFocusedIndex = 0;
 
-        return 0;
-    }, [policyCategories, selectedOptions, isCategoriesCountBelowThreshold]);
+        if (!_.isEmpty(searchValue) || isCategoriesCountBelowThreshold) {
+            const index = _.findIndex(lodashGet(sections, '[0].data', []), (category) => category.searchText === iou.category);
 
-    const sections = useMemo(
-        () => OptionsListUtils.getNewChatOptions({}, {}, [], searchValue, selectedOptions, [], false, false, true, policyCategories, policyRecentlyUsedCategories, false).categoryOptions,
-        [policyCategories, policyRecentlyUsedCategories, searchValue, selectedOptions],
-    );
+            categoryInitialFocusedIndex = index === -1 ? 0 : index;
+        }
+
+        return categoryInitialFocusedIndex;
+    }, [iou.category, searchValue, isCategoriesCountBelowThreshold, sections]);
 
     const headerMessage = OptionsListUtils.getHeaderMessage(lodashGet(sections, '[0].data.length', 0) > 0, false, searchValue);
     const shouldShowTextInput = !isCategoriesCountBelowThreshold;
diff --git a/src/components/OptionsSelector/BaseOptionsSelector.js b/src/components/OptionsSelector/BaseOptionsSelector.js
index bff9f8b6d7d0..aa02701b1c98 100755
--- a/src/components/OptionsSelector/BaseOptionsSelector.js
+++ b/src/components/OptionsSelector/BaseOptionsSelector.js
@@ -107,13 +107,13 @@ class BaseOptionsSelector extends Component {
             });
             return;
         }
-        const newFocusedIndex = this.props.selectedOptions.length;
 
+        const newFocusedIndex = this.props.selectedOptions.length;
         // eslint-disable-next-line react/no-did-update-set-state
         this.setState(
             {
                 allOptions: newOptions,
-                focusedIndex: newFocusedIndex,
+                focusedIndex: _.isNumber(this.props.initialFocusedIndex) ? this.props.initialFocusedIndex : newFocusedIndex,
             },
             () => {
                 // If we just toggled an option on a multi-selection page or cleared the search input, scroll to top