From 579c053270750df29e8097b94dbd1623030ae03e Mon Sep 17 00:00:00 2001
From: Johan Persson <johanopersson@gmail.com>
Date: Wed, 25 Feb 2026 13:36:40 +0100
Subject: [PATCH] Migrate Flex component from useStyles to useDefinition

Signed-off-by: Johan Persson <johanopersson@gmail.com>
---
 packages/ui/src/components/Flex/Flex.tsx      | 47 +++++++------------
 packages/ui/src/components/Flex/definition.ts | 19 +++++---
 packages/ui/src/components/Flex/index.ts      |  2 +-
 packages/ui/src/components/Flex/types.ts      | 15 ++++--
 4 files changed, 40 insertions(+), 43 deletions(-)

diff --git a/packages/ui/src/components/Flex/Flex.tsx b/packages/ui/src/components/Flex/Flex.tsx
index a798a88012..cbaea39c97 100644
--- a/packages/ui/src/components/Flex/Flex.tsx
+++ b/packages/ui/src/components/Flex/Flex.tsx
@@ -15,44 +15,29 @@
  */
 
 import { forwardRef } from 'react';
-import { FlexProps } from './types';
-import clsx from 'clsx';
-import { useStyles } from '../../hooks/useStyles';
+import type { FlexProps } from './types';
+import { useDefinition } from '../../hooks/useDefinition';
 import { FlexDefinition } from './definition';
-import styles from './Flex.module.css';
-import { BgProvider, useBgProvider } from '../../hooks/useBg';
 
 /** @public */
 export const Flex = forwardRef<HTMLDivElement, FlexProps>((props, ref) => {
-  const { bg: resolvedBg } = useBgProvider(props.bg);
+  const { ownProps, restProps, dataAttributes, utilityStyle } = useDefinition(
+    FlexDefinition,
+    { gap: '4', ...props },
+  );
+  const { classes, childrenWithBgProvider } = ownProps;
 
-  const { classNames, dataAttributes, utilityClasses, style, cleanedProps } =
-    useStyles(FlexDefinition, {
-      gap: '4',
-      ...props,
-      bg: resolvedBg, // Use resolved bg for data attribute
-    });
-
-  const { className, bg, ...rest } = cleanedProps;
-
-  const content = (
+  return (
     <div
       ref={ref}
-      className={clsx(
-        classNames.root,
-        utilityClasses,
-        styles[classNames.root],
-        className,
-      )}
-      style={style}
+      className={classes.root}
+      style={{ ...utilityStyle, ...ownProps.style }}
       {...dataAttributes}
-      {...rest}
-    />
-  );
-
-  return resolvedBg ? (
-    <BgProvider bg={resolvedBg}>{content}</BgProvider>
-  ) : (
-    content
+      {...restProps}
+    >
+      {childrenWithBgProvider}
+    </div>
   );
 });
+
+Flex.displayName = 'Flex';
diff --git a/packages/ui/src/components/Flex/definition.ts b/packages/ui/src/components/Flex/definition.ts
index 2701ce5e10..802464c4e1 100644
--- a/packages/ui/src/components/Flex/definition.ts
+++ b/packages/ui/src/components/Flex/definition.ts
@@ -14,16 +14,26 @@
  * limitations under the License.
  */
 
-import type { ComponentDefinition } from '../../types';
+import { defineComponent } from '../../hooks/useDefinition';
+import type { FlexOwnProps } from './types';
+import styles from './Flex.module.css';
 
 /**
  * Component definition for Flex
  * @public
  */
-export const FlexDefinition = {
+export const FlexDefinition = defineComponent<FlexOwnProps>()({
+  styles,
   classNames: {
     root: 'bui-Flex',
   },
+  bg: 'provider',
+  propDefs: {
+    bg: { dataAttribute: true },
+    children: {},
+    className: {},
+    style: {},
+  },
   utilityProps: [
     'm',
     'mb',
@@ -44,7 +54,4 @@ export const FlexDefinition = {
     'justify',
     'direction',
   ],
-  dataAttributes: {
-    bg: ['neutral', 'danger', 'warning', 'success'] as const,
-  },
-} as const satisfies ComponentDefinition;
+});
diff --git a/packages/ui/src/components/Flex/index.ts b/packages/ui/src/components/Flex/index.ts
index c9c5890ad9..244d92c601 100644
--- a/packages/ui/src/components/Flex/index.ts
+++ b/packages/ui/src/components/Flex/index.ts
@@ -15,4 +15,4 @@
  */
 export { Flex } from './Flex';
 export { FlexDefinition } from './definition';
-export type { FlexProps } from './types';
+export type { FlexProps, FlexOwnProps } from './types';
diff --git a/packages/ui/src/components/Flex/types.ts b/packages/ui/src/components/Flex/types.ts
index 1f04a4818d..1e814de3b4 100644
--- a/packages/ui/src/components/Flex/types.ts
+++ b/packages/ui/src/components/Flex/types.ts
@@ -14,16 +14,21 @@
  * limitations under the License.
  */
 
+import type { ReactNode, CSSProperties } from 'react';
 import type { Responsive, Space, SpaceProps, ProviderBg } from '../../types';
 
 /** @public */
-export interface FlexProps extends SpaceProps {
-  children?: React.ReactNode;
+export type FlexOwnProps = {
+  children?: ReactNode;
+  className?: string;
+  style?: CSSProperties;
+  bg?: Responsive<ProviderBg>;
+};
+
+/** @public */
+export interface FlexProps extends SpaceProps, FlexOwnProps {
   gap?: Responsive<Space>;
   align?: Responsive<'start' | 'center' | 'end' | 'baseline' | 'stretch'>;
   justify?: Responsive<'start' | 'center' | 'end' | 'between'>;
   direction?: Responsive<'row' | 'column' | 'row-reverse' | 'column-reverse'>;
-  className?: string;
-  style?: React.CSSProperties;
-  bg?: Responsive<ProviderBg>;
 }