fix: Make Code more like pre/code elements

packages/ui/src/components/Code/Code.stories.tsx

@@ -1,31 +1,20 @@
+import React from 'react';
+
 import { Code } from './Code';
+import { CodeBlock } from './CodeBlock/CodeBlock';
 
 import type { Meta, StoryObj } from '@storybook/react';
 
-const meta: Meta<typeof Code> = {
+const meta: Meta = {
   title: 'Components/Code',
-  component: Code,
   tags: ['autodocs'],
-  argTypes: {
-    variant: {
-      control: 'radio',
-      options: ['block', 'inline'],
-    },
-    showLineNumbers: {
-      control: 'boolean',
-    },
-    label: {
-      control: 'text',
-    },
-  },
 };
 
 export default meta;
-type Story = StoryObj<typeof Code>;
 
 // Helper to simulate highlighted lines
 const MockHighlightedCode = () => (
-  <>
+  <Code>
     <span className="block text-[var(--rb-primary-text)]">
       import React from &apos;react&apos;;
     </span>
@@ -33,73 +22,85 @@ const MockHighlightedCode = () => (
     <span className="block text-[var(--rb-secondary-text)]">
       export const MyComponent = () =&gt; &#123;
     </span>
-    <span className="block text-[var(--rb-text)] pl-4">
-      return &lt;div&gt;Hello World&lt;/div&gt;;
+    <span className="block text-[var(--rb-text)]">
+      {'    '}return &lt;div&gt;Hello World&lt;/div&gt;;
     </span>
     <span className="block text-[var(--rb-secondary-text)]">&#125;;</span>
-  </>
+  </Code>
 );
 
 const LongHighlightedCode = () => (
-  <>
+  <Code>
     {Array.from({ length: 50 }).map((_, i) => (
       <span key={`line-${String(i)}`} className="block">
         <span className="text-[var(--rb-tertiary-text)]">line_{i + 1}</span>:{' '}
         <span className="text-[var(--rb-text)]">const x = {Math.random()};</span>
       </span>
     ))}
-  </>
+  </Code>
 );
 
-export const DefaultBlock: Story = {
+export const DefaultBlock: StoryObj<typeof CodeBlock> = {
+  render: (args) => (
+    <CodeBlock {...args}>
+      <MockHighlightedCode />
+    </CodeBlock>
+  ),
   args: {
-    children: <MockHighlightedCode />,
-    variant: 'block',
     label: 'MyComponent.tsx',
   },
 };
 
-export const WithoutLabel: Story = {
+export const WithoutLabel: StoryObj<typeof CodeBlock> = {
+  render: (args) => (
+    <CodeBlock {...args}>
+      <MockHighlightedCode />
+    </CodeBlock>
+  ),
   args: {
-    children: <MockHighlightedCode />,
-    variant: 'block',
     showLineNumbers: true,
   },
 };
 
-export const WithoutLineNumbers: Story = {
+export const WithoutLineNumbers: StoryObj<typeof CodeBlock> = {
+  render: (args) => (
+    <CodeBlock {...args}>
+      <MockHighlightedCode />
+    </CodeBlock>
+  ),
   args: {
-    children: <MockHighlightedCode />,
-    variant: 'block',
     showLineNumbers: false,
     label: 'no-lines.js',
   },
 };
 
-export const Inline: Story = {
-  args: {
-    children: 'npm install @ermnvldmr/ui',
-    variant: 'inline',
-  },
+export const Inline: StoryObj<typeof Code> = {
+  render: (args) => <Code {...args}>npm install @ermnvldmr/ui</Code>,
 };
 
-export const LongWithScroll: Story = {
+export const LongWithScroll: StoryObj<typeof CodeBlock> = {
+  render: (args) => (
+    <CodeBlock {...args}>
+      <LongHighlightedCode />
+    </CodeBlock>
+  ),
   args: {
-    children: <LongHighlightedCode />,
-    variant: 'block',
     label: 'long-file.ts',
     maxHeight: '300px',
   },
 };
 
-export const PlainText: Story = {
+export const PlainText: StoryObj<typeof CodeBlock> = {
+  render: (args) => (
+    <CodeBlock {...args}>
+      <Code>
+        {`const simple = "string content";
+console.log(simple);`}
+      </Code>
+    </CodeBlock>
+  ),
   args: {
-    children: `const simple = "string content";
-console.log(simple);`,
-    variant: 'block',
     label: 'plain-text.js',
-    // Note: Line numbers won't work automatically on plain text unless we split it into elements
-    // This story demonstrates behavior with raw text (counters will count the single text node as 1 line or 0)
     showLineNumbers: false,
   },
 };

packages/ui/src/components/Code/Code.test.tsx

@@ -2,9 +2,10 @@ import { render, screen, fireEvent, act } from '@testing-library/react';
 import React from 'react';
 
 import { Code } from './Code';
+import { CodeBlock } from './CodeBlock/CodeBlock';
 
 // Mock clipboard API
-const mockWriteText = jest.fn();
+const mockWriteText = jest.fn().mockImplementation(() => Promise.resolve());
 Object.assign(navigator, {
   clipboard: {
     writeText: mockWriteText,
@@ -17,7 +18,7 @@ describe('Code Component', () => {
   });
 
   it('renders inline code correctly', () => {
-    render(<Code variant="inline">const x = 1;</Code>);
+    render(<Code>const x = 1;</Code>);
     const codeElement = screen.getByText('const x = 1;');
     expect(codeElement.tagName).toBe('CODE');
     expect(codeElement).toHaveClass('rounded-sm');
@@ -25,9 +26,11 @@ describe('Code Component', () => {
 
   it('renders block code correctly', () => {
     render(
-      <Code label="test.ts" variant="block">
+      <CodeBlock label="test.ts">
+        <Code>
           <span>console.log(&quot;hello&quot;);</span>
         </Code>
+      </CodeBlock>
     );
     // Should verify it's inside a pre tag
     const preElement = screen.getByText('console.log("hello");').closest('pre');
@@ -37,42 +40,43 @@ describe('Code Component', () => {
     expect(screen.getByText('test.ts')).toBeInTheDocument();
   });
 
-  it('copies content on click in inline mode', () => {
-    render(
-      <Code copyValue="copied!" variant="inline">
-        Display Text
-      </Code>
-    );
+  it('copies content on click in inline mode', async () => {
+    render(<Code copyValue="copied!">Display Text</Code>);
     const codeElement = screen.getByText('Display Text');
 
-    act(() => {
+    await act(async () => {
       fireEvent.click(codeElement);
+      await Promise.resolve();
     });
 
     expect(mockWriteText).toHaveBeenCalledWith('copied!');
   });
 
-  it('copies content via button in block mode', () => {
+  it('copies content via button in block mode', async () => {
     render(
-      <Code copyValue="block copy" variant="block">
+      <CodeBlock copyValue="block copy">
+        <Code>
           <span>some code</span>
         </Code>
+      </CodeBlock>
     );
 
     const copyButton = screen.getByLabelText('Copy code');
-    act(() => {
+    await act(async () => {
       fireEvent.click(copyButton);
+      await Promise.resolve();
     });
 
     expect(mockWriteText).toHaveBeenCalledWith('block copy');
   });
 
-  it('extracts text from children if copyValue is not provided', () => {
-    render(<Code variant="inline">extracted</Code>);
+  it('extracts text from children if copyValue is not provided', async () => {
+    render(<Code>extracted</Code>);
 
     const codeElement = screen.getByText('extracted');
-    act(() => {
+    await act(async () => {
       fireEvent.click(codeElement);
+      await Promise.resolve();
     });
 
     expect(mockWriteText).toHaveBeenCalledWith('extracted');

packages/ui/src/components/Code/Code.tsx

@@ -1,69 +1,39 @@
 import { cn } from '@ermnvldmr/stl';
-import { Check, Copy } from 'lucide-react';
 import React, { memo, useCallback, useEffect, useRef, useState } from 'react';
 import { usePress } from 'react-aria';
 
-import { Button } from '../Button/Button';
-import { Container } from '../Container/Container';
+import { useCodeContext } from './contexts/CodeContext/CodeContext';
 
 import type { ClassNameProps, TestIdProps } from '@ermnvldmr/stl';
 
 /**
- *
+ * Props for the Code component.
  */
 export interface CodeProps extends ClassNameProps, TestIdProps {
   /** The content to display. Can be a string or pre-highlighted HTML elements. */
   children: React.ReactNode;
 
   /**
-   * Display style of the code.
-   * - 'block': Multi-line container with optional scroll, line numbers, and header controls.
-   * - 'inline': Compact, clickable snippet for embedding in text.
-   * @default 'block'
-   */
-  variant?: 'block' | 'inline';
-
-  /**
-   * Optional label text (e.g., filename "index.ts", language "bash")
-   * displayed in the top-left corner of block variants.
-   */
-  label?: string;
-
-  /**
-   * Whether to render line number gutters in block mode.
-   * Requires children to be rendered as block-level elements (e.g., via a highlighter).
-   * @default true
-   */
-  showLineNumbers?: boolean;
-
-  /**
-   * Raw string value to copy to clipboard.
-   * If omitted, the component will attempt to extract text from `children`.
+   * Raw string value to copy to clipboard (only used in inline mode).
+   * If omitted in inline mode, the component will attempt to extract text from `children`.
    */
   copyValue?: string;
-
-  /**
-   * Maximum height of the code block before vertical scrolling is enabled.
-   */
-  maxHeight?: string | number;
 }
 
 /**
- * A general-purpose wrapper for displaying source code.
+ * A component for displaying source code.
  *
- * It supports both block-level (multi-line) and inline (snippet) display modes,
- * with features like line numbering, syntax highlighting containers, and copy-to-clipboard functionality.
+ * When used inside a `CodeBlock`, it renders a block-level `<code>` element
+ * with support for line numbering.
+ * When used standalone, it renders an inline snippet that is clickable to copy.
  */
 export const Code = memo(function Code({
   children,
-  variant = 'block',
-  label,
-  showLineNumbers = true,
   copyValue,
-  maxHeight,
   className,
   'data-testid': testId,
 }: CodeProps) {
+  const context = useCodeContext();
   const [isCopied, setIsCopied] = useState(false);
   const codeRef = useRef<HTMLElement>(null);
 
@@ -97,12 +67,42 @@ export const Code = memo(function Code({
     }
   }, [isCopied]);
 
-  // Inline Variant
   const { pressProps } = usePress({
     onPress: onCopyPress,
   });
 
-  if (variant === 'inline') {
+  if (context?.isInCodeBlock) {
+    const { showLineNumbers } = context;
+    return (
+      <code
+        ref={codeRef}
+        className={cn(
+          'block min-w-full font-mono whitespace-pre',
+          showLineNumbers && [
+            '[counter-reset:line]',
+            '[&>*]:relative [&>*]:block [&>*]:pl-10',
+            '[&>*]:[counter-increment:line]',
+            '[&>*]:before:content-[counter(line)]',
+            '[&>*]:before:absolute [&>*]:before:left-0',
+            '[&>*]:before:inline-block',
+            '[&>*]:before:w-8',
+            '[&>*]:before:text-right',
+            '[&>*]:before:text-[var(--rb-muted-text)]/50',
+            '[&>*]:before:select-none',
+            '[&>*]:before:border-r',
+            '[&>*]:before:border-[var(--rb-outline)]/10',
+            '[&>*]:before:pr-3',
+            '[&>*]:before:font-mono [&>*]:before:text-[10px]',
+          ],
+          className
+        )}
+        data-testid={testId}
+      >
+        {children}
+      </code>
+    );
+  }
+
   return (
     <code
       {...pressProps}
@@ -124,78 +124,4 @@ export const Code = memo(function Code({
       {children}
     </code>
   );
-  }
-
-  // Block Variant
-  return (
-    <Container
-      border
-      bg="base"
-      className={cn('relative group overflow-hidden', className)}
-      data-testid={testId}
-      rounded="md"
-    >
-      {/* Header Controls */}
-      <div className="absolute top-0 left-0 right-0 h-0 z-10 pointer-events-none">
-        {label && (
-          <span className="absolute top-2 left-4 text-xs font-mono text-[var(--rb-muted-text)] select-none bg-[var(--rb-container-base)]/80 backdrop-blur-sm px-1.5 rounded pointer-events-auto">
-            {label}
-          </span>
-        )}
-        <div className="absolute top-2 right-2 pointer-events-auto">
-          <Button
-            aria-label="Copy code"
-            className={cn(
-              'h-7 w-7 p-0 bg-[var(--rb-container-base)]/80 backdrop-blur-sm transition-opacity',
-              !isCopied && 'opacity-0 group-hover:opacity-100 focus:opacity-100'
-            )}
-            color="neutral"
-            size="s"
-            variant="ghost"
-            onPress={onCopyPress}
-          >
-            {isCopied ? (
-              <Check className="w-3.5 h-3.5 text-[var(--rb-secondary-text)]" />
-            ) : (
-              <Copy className="w-3.5 h-3.5" />
-            )}
-          </Button>
-        </div>
-      </div>
-
-      {/* Code Area */}
-      <pre
-        className={cn(
-          'overflow-auto p-4 text-sm font-mono leading-relaxed select-text',
-          label && 'pt-8' // Add padding if label is present to avoid overlap
-        )}
-        style={{ maxHeight: maxHeight }}
-      >
-        <code
-          ref={codeRef}
-          className={cn(
-            'block min-w-full',
-            showLineNumbers && [
-              '[counter-reset:line] pl-10',
-              '[&>*]:relative [&>*]:block',
-              '[&>*]:[counter-increment:line]',
-              '[&>*]:before:content-[counter(line)]',
-              '[&>*]:before:absolute [&>*]:before:left-[-2.5rem]',
-              '[&>*]:before:inline-block',
-              '[&>*]:before:w-8',
-              '[&>*]:before:text-right',
-              '[&>*]:before:text-[var(--rb-muted-text)]/50',
-              '[&>*]:before:select-none',
-              '[&>*]:before:border-r',
-              '[&>*]:before:border-[var(--rb-outline)]/10',
-              '[&>*]:before:pr-3',
-              '[&>*]:before:font-mono [&>*]:before:text-[10px]', // Use mono for numbers too
-            ]
-          )}
-        >
-          {children}
-        </code>
-      </pre>
-    </Container>
-  );
 });

packages/ui/src/components/Code/CodeBlock/CodeBlock.tsx

@@ -0,0 +1,148 @@
+import { cn } from '@ermnvldmr/stl';
+import { Check, Copy } from 'lucide-react';
+import React, { memo, useCallback, useEffect, useRef, useState, useMemo } from 'react';
+
+import { Button } from '../../Button/Button';
+import { Container } from '../../Container/Container';
+import { CodeContext, type CodeContextValue } from '../contexts/CodeContext/CodeContext';
+
+import type { ClassNameProps, TestIdProps } from '@ermnvldmr/stl';
+
+/**
+ * Props for the CodeBlock component.
+ */
+export interface CodeBlockProps extends ClassNameProps, TestIdProps {
+  /** The content to display, typically a Code component. */
+  children: React.ReactNode;
+
+  /**
+   * Optional label text (e.g., filename "index.ts", language "bash")
+   * displayed in the top-left corner.
+   */
+  label?: string;
+
+  /**
+   * Whether to render line number gutters.
+   * Requires children (Code) to be rendered as block-level elements.
+   * @default true
+   */
+  showLineNumbers?: boolean;
+
+  /**
+   * Raw string value to copy to clipboard.
+   * If omitted, the component will attempt to extract text from its content.
+   */
+  copyValue?: string;
+
+  /**
+   * Maximum height of the code block before vertical scrolling is enabled.
+   */
+  maxHeight?: string | number;
+}
+
+/**
+ * A container for multi-line source code display.
+ *
+ * It provides a framed container with optional label, copy-to-clipboard button,
+ * and line numbering support for its children.
+ */
+export const CodeBlock = memo(function CodeBlock({
+  children,
+  label,
+  showLineNumbers = true,
+  copyValue,
+  maxHeight,
+  className,
+  'data-testid': testId,
+}: CodeBlockProps) {
+  const [isCopied, setIsCopied] = useState(false);
+  const preRef = useRef<HTMLPreElement>(null);
+
+  const handleCopy = useCallback(async () => {
+    try {
+      let textToCopy = copyValue;
+      if (!textToCopy && preRef.current) {
+        textToCopy = preRef.current.innerText;
+        if (!textToCopy) {
+          textToCopy = preRef.current.textContent ?? undefined;
+        }
+      }
+
+      if (textToCopy) {
+        await navigator.clipboard.writeText(textToCopy);
+        setIsCopied(true);
+      }
+    } catch (err) {
+      console.error('Failed to copy code:', err);
+    }
+  }, [copyValue]);
+
+  const onCopyPress = useCallback(() => {
+    void handleCopy();
+  }, [handleCopy]);
+
+  useEffect(() => {
+    if (isCopied) {
+      const timer = setTimeout(() => setIsCopied(false), 2000);
+      return () => clearTimeout(timer);
+    }
+  }, [isCopied]);
+
+  const contextValue = useMemo<CodeContextValue>(
+    () => ({
+      isInCodeBlock: true,
+      showLineNumbers,
+    }),
+    [showLineNumbers]
+  );
+
+  return (
+    <Container
+      border
+      bg="base"
+      className={cn('relative group overflow-hidden', className)}
+      data-testid={testId}
+      rounded="md"
+    >
+      {/* Header Controls */}
+      <div className="absolute top-0 left-0 right-0 h-0 z-10 pointer-events-none">
+        {label && (
+          <span className="absolute top-2 left-4 text-xs font-mono text-[var(--rb-muted-text)] select-none bg-[var(--rb-container-base)]/80 backdrop-blur-sm px-1.5 rounded pointer-events-auto">
+            {label}
+          </span>
+        )}
+        <div className="absolute top-2 right-2 pointer-events-auto">
+          <Button
+            aria-label="Copy code"
+            className={cn(
+              'h-7 w-7 p-0 bg-[var(--rb-container-base)]/80 backdrop-blur-sm transition-opacity',
+              !isCopied && 'opacity-0 group-hover:opacity-100 focus:opacity-100'
+            )}
+            color="neutral"
+            size="s"
+            variant="ghost"
+            onPress={onCopyPress}
+          >
+            {isCopied ? (
+              <Check className="w-3.5 h-3.5 text-[var(--rb-secondary-text)]" />
+            ) : (
+              <Copy className="w-3.5 h-3.5" />
+            )}
+          </Button>
+        </div>
+      </div>
+
+      {/* Code Area */}
+      <pre
+        ref={preRef}
+        className={cn(
+          'overflow-auto p-4 text-sm font-mono leading-relaxed select-text whitespace-pre',
+          label && 'pt-8'
+        )}
+        style={{ maxHeight: maxHeight }}
+      >
+        <CodeContext.Provider value={contextValue}>{children}</CodeContext.Provider>
+      </pre>
+    </Container>
+  );
+});

packages/ui/src/components/Code/contexts/CodeContext/CodeContext.ts

@@ -0,0 +1,33 @@
+import { createContext, useContext } from 'react';
+
+/**
+ * Value for the Code context.
+ */
+export interface CodeContextValue {
+  /** Whether the code is rendered inside a CodeBlock. */
+  isInCodeBlock: boolean;
+  /** Whether to show line numbers in the block. */
+  showLineNumbers?: boolean;
+}
+
+/**
+ * Context for sharing code-related state between CodeBlock and Code components.
+ */
+export const CodeContext = createContext<CodeContextValue | null>(null);
+
+/**
+ * Hook to access the Code context.
+ *
+ * @returns The current code context value or null if not within a CodeBlock.
+ *
+ * @example
+ * ```tsx
+ * const context = useCodeContext();
+ * if (context?.isInCodeBlock) {
+ *   // ...
+ * }
+ * ```
+ */
+export function useCodeContext(): CodeContextValue | null {
+  return useContext(CodeContext);
+}

packages/ui/src/components/Code/index.ts

@@ -0,0 +1,2 @@
+export * from './Code';
+export * from './CodeBlock/CodeBlock';

packages/ui/src/components/index.ts

@@ -10,7 +10,7 @@ export * from './List/List';
 export * from './Markdown/Markdown';
 export * from './Paragraph/Paragraph';
 export * from './Button/Button';
-export * from './Code/Code';
+export * from './Code';
 export * from './Blockquote';
 export * from './Stub/Stub';
 export * from './Image/Image';
@@ -29,7 +29,7 @@ export type { TextProps, TextType, TextSize, TextColor } from './Text/Text';
 export type { HeaderProps } from './Header/Header';
 export type { ParagraphProps } from './Paragraph/Paragraph';
 export type { ButtonProps, ButtonVariant, ButtonColor, ButtonSize } from './Button/Button';
-export type { CodeProps } from './Code/Code';
+export type { CodeProps, CodeBlockProps } from './Code';
 export type { BlockquoteProps, BlockquoteCitationProps } from './Blockquote';
 export type { StubProps } from './Stub/Stub';
 export type { ImageProps } from './Image/Image';