fix bug and address comments

pages/bluedialog/bluedialog.page.tsx

@@ -1,6 +1,6 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-import React, { useRef } from 'react';
+import React, { useEffect, useRef } from 'react';
 
 import { Box, Button, Checkbox, Form, FormField, SpaceBetween, Textarea } from '~components';
 
@@ -20,21 +20,24 @@ export default function GenAIFeedback() {
   });
   const [feedbackText, setFeedbackText] = React.useState('');
   const checkboxRef = useRef<HTMLElement>(null);
+  const dialogRef = useRef<HTMLDivElement>(null);
+
+  useEffect(() => {
+    dialogRef.current?.focus();
+  }, [dialogRef]);
 
   const selectOption = (option: string, checked: boolean) => {
     setFeedbackOptions({ ...feedbackOptions, [option]: checked });
     setErrorMessage(null);
   };
 
-  const submitData = () => {
+  const submitFeedback = () => {
     // Validation
     const isFeedbackOptionSelected = Object.values(feedbackOptions).some(val => !!val);
     if (!isFeedbackOptionSelected) {
       setErrorMessage('At least one option must be selected.');
       // Move focus to the required input
-      if (checkboxRef.current) {
-        checkboxRef.current.focus();
-      }
+      checkboxRef.current?.focus();
       return;
     }
 
@@ -51,10 +54,11 @@ export default function GenAIFeedback() {
     <Box padding="xxl">
       {showDialog && (
         <FeedbackDialog
+          ref={dialogRef}
           onDismiss={dismissDialog}
           footer={
             <div className={styles['footer-buttons']}>
-              <Button onClick={submitData} ariaLabel="Submit form">
+              <Button onClick={submitFeedback} ariaLabel="Submit form">
                 Submit
               </Button>
             </div>

pages/bluedialog/dialog.tsx

@@ -1,57 +1,45 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-import React, { useEffect, useRef } from 'react';
+import React from 'react';
 
 import { Button } from '~components';
 
 import styles from './styles.scss';
 
-export default function FeedbackDialog({
-  children,
-  footer,
-  onDismiss,
-}: {
-  children: React.ReactNode;
-  footer: React.ReactNode;
-  onDismiss: () => void;
-}) {
-  const dialogRef = useRef<HTMLDivElement>(null);
-  const triggerRef = useRef<HTMLElement | null>(null);
-
-  useEffect(() => {
-    // Store the element that had focus before dialog opened
-    triggerRef.current = document.activeElement as HTMLElement;
-
-    // Focus the dialog container when it opens
-    if (dialogRef.current) {
-      dialogRef.current.focus();
-    }
-    // Cleanup: Return focus to the trigger element when dialog closes
-    return () => {
-      if (triggerRef.current) {
-        triggerRef.current.focus();
-      }
-    };
-  }, []);
-
-  return (
-    <div
-      ref={dialogRef}
-      className={styles['feedback-dialog']}
-      role="dialog"
-      aria-label="Feedback dialog"
-      aria-modal="false" // Maintains natural focus flow since it's an inline dialog
-      tabIndex={-1} // This allows the dialog to receive focus
-    >
-      <div className={styles.content}>
-        <div className={styles.dismiss}>
-          <Button iconName="close" variant="icon" onClick={onDismiss} ariaLabel="Close dialog" />
+const FeedbackDialog = React.forwardRef(
+  (
+    {
+      children,
+      footer,
+      onDismiss,
+    }: {
+      children: React.ReactNode;
+      footer: React.ReactNode;
+      onDismiss: () => void;
+    },
+    ref: React.Ref<HTMLDivElement>
+  ) => {
+    return (
+      <div
+        ref={ref}
+        className={styles['feedback-dialog']}
+        role="dialog"
+        aria-label="Feedback dialog"
+        aria-modal="false" // Maintains natural focus flow since it's an inline dialog
+        tabIndex={-1} // This allows the dialog to receive focus
+      >
+        <div className={styles.content}>
+          <div className={styles.dismiss}>
+            <Button iconName="close" variant="icon" onClick={onDismiss} ariaLabel="Close dialog" />
+          </div>
+
+          <div className={styles['inner-content']}>{children}</div>
         </div>
 
-        <div className={styles['inner-content']}>{children}</div>
+        <div className={styles.footer}>{footer}</div>
       </div>
+    );
+  }
+);
 
-      <div className={styles.footer}>{footer}</div>
-    </div>
-  );
-}
+export default FeedbackDialog;

pages/bluedialog/styles.scss

@@ -21,18 +21,20 @@
 .content {
   display: flex;
   align-items: baseline;
-
-  padding-block-start: awsui.$space-scaled-s;
-  padding-block-end: awsui.$space-scaled-l;
-  padding-inline: awsui.$space-container-horizontal;
 }
 
 .dismiss {
   order: 1;
+  padding-inline-end: awsui.$space-container-horizontal;
 }
 
 .inner-content {
   flex: 1;
+  overflow: hidden; // prevent textarea from growing outside of the container
+
+  padding-block-start: awsui.$space-scaled-s;
+  padding-block-end: awsui.$space-scaled-l;
+  padding-inline-start: awsui.$space-container-horizontal;
 }
 
 .footer {