From eb4a0492df7e414f311de7f4d6c553837ad14453 Mon Sep 17 00:00:00 2001
From: Myrto Kontouli <myrto.kontouli@atomlearning.co.uk>
Date: Tue, 25 Jul 2023 20:03:26 +0100
Subject: [PATCH] feat: update Flex to do all that Stack does

---
 .../editor-component/ColorTokenList.tsx       |  4 +-
 lib/src/components/flex/Flex.test.tsx         | 12 ++-
 lib/src/components/flex/Flex.tsx              | 95 ++++++++++++++++++-
 .../flex/__snapshots__/Flex.test.tsx.snap     | 14 +--
 4 files changed, 110 insertions(+), 15 deletions(-)

diff --git a/documentation/components/markdown/editor-component/ColorTokenList.tsx b/documentation/components/markdown/editor-component/ColorTokenList.tsx
index 704fe2136..4ec5b4704 100644
--- a/documentation/components/markdown/editor-component/ColorTokenList.tsx
+++ b/documentation/components/markdown/editor-component/ColorTokenList.tsx
@@ -10,9 +10,9 @@ const ColorExample: typeof TokenList.Item = ({ token, value, ...rest }) => {
 
   const hasAlpha = color.hasOwnProperty('alpha')
   return (
-    <Flex css={{ alignItems: 'center' }} {...rest}>
+    <Flex align="center" {...rest}>
       <Box css={{ borderRadius: '$round', bg: `$${token}`, size: '$6' }} />
-      <Flex css={{ pl: '$3', flexDirection: 'column' }}>
+      <Flex direction="column" css={{ pl: '$3' }}>
         <Text css={{ fontWeight: 600, mb: '$3' }}>{`$${token}`}</Text>
         <Text size="sm" css={{ color: '$base8', mb: !hasAlpha ? '$3' : 0 }}>
           {value}
diff --git a/lib/src/components/flex/Flex.test.tsx b/lib/src/components/flex/Flex.test.tsx
index b1defdf25..9de84171c 100644
--- a/lib/src/components/flex/Flex.test.tsx
+++ b/lib/src/components/flex/Flex.test.tsx
@@ -4,17 +4,19 @@ import * as React from 'react'
 
 import { Flex } from '.'
 
+const FlexImplementation = (props) => {
+  return <Flex align="center" justify="center" {...props} />
+}
+
 describe(`Flex component`, () => {
   it('renders', async () => {
-    const { container } = render(
-      <Flex css={{ m: 'auto', height: 100, width: 100 }}>FLEX</Flex>
-    )
-    await screen.findByText('FLEX')
+    const { container } = render(<FlexImplementation>Flex</FlexImplementation>)
+    await screen.findByText('Flex')
     expect(container).toMatchSnapshot()
   })
 
   it('has no programmatically detectable a11y issues', async () => {
-    render(<Flex />)
+    render(<FlexImplementation />)
 
     const results = await axe(document.body)
     expect(results).toHaveNoViolations()
diff --git a/lib/src/components/flex/Flex.tsx b/lib/src/components/flex/Flex.tsx
index 528a3681f..77384d174 100644
--- a/lib/src/components/flex/Flex.tsx
+++ b/lib/src/components/flex/Flex.tsx
@@ -1,5 +1,96 @@
-import { styled } from '~/stitches'
+import { CSS, styled } from '~/stitches'
+import { createThemeVariants } from '~/utilities'
 
-export const Flex = styled('div', { display: 'flex' })
+const createVariants = <T extends GlobalValue[]>(
+  variants: T,
+  fn: (string: T[number]) => CSS
+) => {
+  return variants.reduce(
+    (prev, variant) => ({ ...prev, [variant]: fn(variant) }),
+    {} as Record<T[number], CSS>
+  )
+}
+
+const globalValues = [
+  'inherit',
+  'initial',
+  'revert',
+  'revert-layer',
+  'unset'
+] as const
+
+/*
+ * The following type is partially a hack to get the rest of the createVariants array parameter to be recognised as const.
+ * Thus giving is the correct types generated for these generated variant props.
+ * No clue how we can do it cleaner.
+ * Mad props to: Elliot for getting this to working as is.
+ *
+ */
+// eslint-disable-next-line @typescript-eslint/ban-types
+type GlobalValue = typeof globalValues[number] | (string & {})
+
+export const Flex = styled('div', {
+  display: 'flex',
+  variants: {
+    direction: createVariants(
+      [...globalValues, 'row', 'row-reverse', 'column', 'column-reverse'],
+      (v) => {
+        return { flexDirection: v }
+      }
+    ),
+    wrap: createVariants(
+      [...globalValues, 'nowrap', 'wrap', 'wrap-reverse'],
+      (v) => {
+        return { flexWrap: v }
+      }
+    ),
+    // Why is both `start` and `flex-start` valid values? Answer: https://csslayout.news/whats-the-difference-between-the-alignment-values-of-start-flex-start-and-self-start/
+    justify: createVariants(
+      [
+        ...globalValues,
+        'normal',
+        'unsafe',
+        'safe',
+        'start',
+        'center',
+        'end',
+        'flex-start',
+        'flex-end',
+        'left',
+        'right',
+        'space-between',
+        'space-around',
+        'space-evenly',
+        'stretch'
+      ],
+      (v) => {
+        return { justifyContent: v }
+      }
+    ),
+    align: createVariants(
+      [
+        ...globalValues,
+        'normal',
+        'unsafe',
+        'safe',
+        'center',
+        'start',
+        'end',
+        'self-start',
+        'self-end',
+        'flex-start',
+        'flex-end',
+        'baseline',
+        'first baseline',
+        'last baseline',
+        'stretch'
+      ],
+      (v) => {
+        return { alignItems: v }
+      }
+    ),
+    gap: createThemeVariants('space', { gap: '$key' })
+  }
+})
 
 Flex.displayName = 'Flex'
diff --git a/lib/src/components/flex/__snapshots__/Flex.test.tsx.snap b/lib/src/components/flex/__snapshots__/Flex.test.tsx.snap
index 85bad1ed4..405ed8adb 100644
--- a/lib/src/components/flex/__snapshots__/Flex.test.tsx.snap
+++ b/lib/src/components/flex/__snapshots__/Flex.test.tsx.snap
@@ -8,18 +8,20 @@ exports[`Flex component renders 1`] = `
 }
 
 @media  {
-  .c-dhzjXW-ikLAKdM-css {
-    margin: auto;
-    height: 100px;
-    width: 100px;
+  .c-dhzjXW-bICGYT-justify-center {
+    justify-content: center;
+  }
+
+  .c-dhzjXW-jroWjL-align-center {
+    align-items: center;
   }
 }
 
 <div>
   <div
-    class="c-dhzjXW c-dhzjXW-ikLAKdM-css"
+    class="c-dhzjXW c-dhzjXW-bICGYT-justify-center c-dhzjXW-jroWjL-align-center"
   >
-    FLEX
+    Flex
   </div>
 </div>
 `;