From 33a4bf6fa7af7d9989e9df157ddee66ea27cc10f Mon Sep 17 00:00:00 2001
From: Rafael Bardini <rbardini@users.noreply.github.com>
Date: Sun, 13 Jun 2021 16:44:16 +0200
Subject: [PATCH] feat(mode): support docs view mode (#27)

---
 example/.storybook/main.js |  6 +++-
 example/package.json       |  1 +
 src/index.tsx              | 69 +++++++++++++++++++++++++++++---------
 src/register.tsx           |  2 +-
 4 files changed, 60 insertions(+), 18 deletions(-)

diff --git a/example/.storybook/main.js b/example/.storybook/main.js
index fd1c803..5a50e0c 100644
--- a/example/.storybook/main.js
+++ b/example/.storybook/main.js
@@ -1,4 +1,8 @@
 module.exports = {
-  addons: ['@storybook/addon-storysource', 'storybook-addon-paddings'],
+  addons: [
+    '@storybook/addon-docs',
+    '@storybook/addon-storysource',
+    'storybook-addon-paddings',
+  ],
   stories: ['../src/**/*.stories.js'],
 };
diff --git a/example/package.json b/example/package.json
index a8ca3a6..af0533f 100644
--- a/example/package.json
+++ b/example/package.json
@@ -11,6 +11,7 @@
   },
   "devDependencies": {
     "@babel/core": "^7.0.0",
+    "@storybook/addon-docs": "^6.2.0",
     "@storybook/addon-storysource": "^6.2.0",
     "@storybook/addons": "^6.2.0",
     "@storybook/react": "^6.2.0",
diff --git a/src/index.tsx b/src/index.tsx
index f6eb49c..3bb45ea 100644
--- a/src/index.tsx
+++ b/src/index.tsx
@@ -1,29 +1,66 @@
-import { StoryWrapper, useMemo } from '@storybook/addons';
+import { StoryWrapper, useEffect, useMemo } from '@storybook/addons';
 
 import { DEFAULT_PADDING, PARAM_KEY } from './constants';
 import { normalizeValues, getSelectedPadding } from './helpers';
 
-const setStyle = (padding = DEFAULT_PADDING) => {
-  document.body.style.margin = '0';
-  document.body.style.padding = padding;
-  document.body.style.transition = 'padding .3s';
+const setStyle = (selector: string, css: string) => {
+  const existingStyle = document.getElementById(selector);
+
+  if (existingStyle) {
+    if (existingStyle.innerHTML !== css) {
+      existingStyle.innerHTML = css;
+    }
+  } else {
+    const style = document.createElement('style');
+    style.setAttribute('id', selector);
+    style.innerHTML = css;
+
+    document.head.appendChild(style);
+  }
 };
 
-export const withPaddings: StoryWrapper = (getStory, context) => {
-  const { globals, parameters } = context;
+const WithPaddings: StoryWrapper = (getStory, context) => {
+  const { id, globals, parameters, viewMode } = context;
   const globalsSelectedPadding = globals[PARAM_KEY]?.value;
   const paddingsConfig = parameters[PARAM_KEY];
+  const isInDocs = viewMode === 'docs';
+  const selector = isInDocs
+    ? `#anchor--${id} .docs-story > div`
+    : '.sb-show-main';
+
+  const selectedPadding = useMemo(
+    () =>
+      getSelectedPadding(
+        normalizeValues(paddingsConfig),
+        globalsSelectedPadding,
+      ) ?? DEFAULT_PADDING,
+    [paddingsConfig, globalsSelectedPadding],
+  );
+
+  const paddingStyles = useMemo(
+    () => `
+      ${selector} {
+        margin: 0;
+        padding: ${selectedPadding} !important;
+        transition: padding .3s;
+      }
+
+      ${selector} .innerZoomElementWrapper > div {
+        border-width: 0 !important;
+      }
+    `,
+    [selector, selectedPadding],
+  );
+
+  useEffect(() => {
+    const selectorId = isInDocs
+      ? `addon-paddings-docs-${id}`
+      : `addon-paddings`;
 
-  // eslint-disable-next-line react-hooks/rules-of-hooks
-  useMemo(() => {
-    const selectedPadding = getSelectedPadding(
-      normalizeValues(paddingsConfig),
-      globalsSelectedPadding,
-    );
-    setStyle(selectedPadding);
-  }, [paddingsConfig, globalsSelectedPadding]);
+    setStyle(selectorId, paddingStyles);
+  }, [id, isInDocs, paddingStyles]);
 
   return getStory(context);
 };
 
-export default withPaddings;
+export default WithPaddings;
diff --git a/src/register.tsx b/src/register.tsx
index cf839bc..8b74fcf 100644
--- a/src/register.tsx
+++ b/src/register.tsx
@@ -8,7 +8,7 @@ addons.register(ADDON_ID, () => {
   addons.add(ADDON_ID, {
     title: 'Paddings',
     type: types.TOOL,
-    match: ({ viewMode }) => viewMode === 'story',
+    match: ({ viewMode }) => !!(viewMode && viewMode.match(/^(story|docs)$/)),
     render: () => <PaddingSelector />,
     paramKey: PARAM_KEY,
   });