dss/dss/storybook/generator.py

"""
Storybook Story Generator for Design System Components

Generates interactive Storybook stories for design system components,
creating comprehensive documentation that showcases component usage,
variants, and integration points.

Stories serve as the primary documentation and interactive reference
for how components should be used in applications.

Supports both React components and Web Components with JSDoc annotations.
"""

import re
import json
import subprocess
import logging
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass, field
from enum import Enum

log = logging.getLogger(__name__)


class StoryTemplate(str, Enum):
    """
    Available story format templates for component documentation.
    """
    CSF3 = "csf3"           # Component Story Format 3 (latest, recommended)
    CSF2 = "csf2"           # Component Story Format 2 (legacy)
    MDX = "mdx"             # MDX format (documentation + interactive)


class ComponentType(str, Enum):
    """
    Type of component detected during parsing.
    """
    REACT = "react"         # React functional or class component
    WEB_COMPONENT = "web"   # Custom Element / Web Component
    UNKNOWN = "unknown"


@dataclass
class PropInfo:
    """
    Component property metadata.

    Captures prop name, type, required status, default value,
    description, and valid options for code generation.
    """
    name: str
    type: str = "unknown"
    required: bool = False
    default_value: Optional[str] = None
    description: str = ""
    options: List[str] = field(default_factory=list)  # For enum/union types


@dataclass
class ComponentMeta:
    """
    Component metadata for story generation.

    Describes component name, file path, props, description,
    and whether it accepts children for story creation.
    """
    name: str
    path: str
    props: List[PropInfo] = field(default_factory=list)
    description: str = ""
    has_children: bool = False
    component_type: ComponentType = ComponentType.UNKNOWN
    tag_name: Optional[str] = None  # For Web Components (e.g., 'ds-button')


class StoryGenerator:
    """
    Story generator for design system components.

    Generates interactive Storybook stories in CSF3, CSF2, or MDX format,
    automatically extracting component metadata and creating comprehensive
    documentation with variants and default stories.

    Supports both React components (.tsx, .jsx) and Web Components (.js).
    """

    def __init__(self, root_path: str):
        self.root = Path(root_path).resolve()
        # Path to the Babel parser script (shared with project_analyzer)
        self._parser_script = Path(__file__).parent.parent / 'analyze' / 'parser.js'

    def generate(self, template: StoryTemplate = StoryTemplate.CSF3, dry_run: bool = True) -> List[Dict[str, str]]:
        """
        Generate stories for all components in the project.

        This is the main entry point for story generation, scanning common
        component directories and generating stories for each component found.

        Args:
            template: Story template format (CSF3, CSF2, or MDX)
            dry_run: If True, only return what would be generated without writing files

        Returns:
            List of dicts with component paths and generated stories
        """
        import asyncio

        # Common component directories to scan
        component_dirs = [
            'src/components',
            'components',
            'src/ui',
            'ui',
            'lib/components',
            'packages/ui/src',
            'app/components',
        ]

        results = []

        for dir_path in component_dirs:
            full_path = self.root / dir_path
            if full_path.exists():
                # Run async method synchronously
                try:
                    loop = asyncio.get_event_loop()
                except RuntimeError:
                    loop = asyncio.new_event_loop()
                    asyncio.set_event_loop(loop)

                dir_results = loop.run_until_complete(
                    self.generate_stories_for_directory(dir_path, template, dry_run)
                )
                results.extend(dir_results)

        # If no component directories found, try root
        if not results:
            try:
                loop = asyncio.get_event_loop()
            except RuntimeError:
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)

            results = loop.run_until_complete(
                self.generate_stories_for_directory('.', template, dry_run)
            )

        return results

    async def generate_story(
        self,
        component_path: str,
        template: StoryTemplate = StoryTemplate.CSF3,
        include_variants: bool = True,
        output_path: Optional[str] = None,
    ) -> str:
        """
        Generate a Storybook story for a component.

        Args:
            component_path: Path to the component file
            template: Story template format
            include_variants: Generate variant stories
            output_path: Optional path to write the story file

        Returns:
            Generated story code
        """
        # Parse component
        meta = await self._parse_component(component_path)

        # Generate story based on template
        if template == StoryTemplate.CSF3:
            story = self._generate_csf3(meta, include_variants)
        elif template == StoryTemplate.CSF2:
            story = self._generate_csf2(meta, include_variants)
        else:
            story = self._generate_mdx(meta, include_variants)

        # Write to file if output path provided
        if output_path:
            output = Path(output_path)
            output.parent.mkdir(parents=True, exist_ok=True)
            output.write_text(story)

        return story

    async def _parse_component(self, component_path: str) -> ComponentMeta:
        """
        Parse a component to extract metadata using Babel AST.

        Supports both React components (TypeScript interfaces) and
        Web Components (JSDoc annotations, observedAttributes, customElements.define).
        """
        path = self.root / component_path if not Path(component_path).is_absolute() else Path(component_path)
        content = path.read_text(encoding="utf-8", errors="ignore")
        component_name = path.stem

        # Try to determine component type and parse accordingly
        component_type = ComponentType.UNKNOWN
        tag_name = None
        props = []
        description = ""
        has_children = False

        # Detect Web Component patterns
        is_web_component = (
            'extends HTMLElement' in content or
            'customElements.define' in content
        )

        if is_web_component:
            component_type = ComponentType.WEB_COMPONENT
            props, tag_name, description, has_children = self._parse_web_component(content)
            # Convert tag name to PascalCase for story naming
            if tag_name:
                component_name = self._tag_to_pascal(tag_name)
        else:
            # Try React component parsing (TypeScript interfaces)
            component_type = ComponentType.REACT
            props = self._parse_react_props(content)

            # Check if component uses children (React)
            has_children = 'children' in content.lower() and (
                'React.ReactNode' in content or
                'ReactNode' in content or
                '{children}' in content
            )

        # Extract description from JSDoc if not already found
        if not description:
            jsdoc_match = re.search(r'/\*\*\s*\n\s*\*\s*([^\n*]+)', content)
            if jsdoc_match:
                description = jsdoc_match.group(1).strip()

        return ComponentMeta(
            name=component_name,
            path=component_path,
            props=props,
            description=description,
            has_children=has_children,
            component_type=component_type,
            tag_name=tag_name,
        )

    def _parse_web_component(self, content: str) -> Tuple[List[PropInfo], Optional[str], str, bool]:
        """
        Parse Web Component to extract attributes from JSDoc and observedAttributes.

        Returns:
            Tuple of (props, tag_name, description, has_children)
        """
        props = []
        tag_name = None
        description = ""
        has_children = False

        # Extract tag name from customElements.define('tag-name', ClassName)
        define_match = re.search(
            r"customElements\.define\s*\(\s*['\"]([^'\"]+)['\"]",
            content
        )
        if define_match:
            tag_name = define_match.group(1)

        # Extract observedAttributes
        # static get observedAttributes() { return ['variant', 'size', ...]; }
        observed_match = re.search(
            r"static\s+get\s+observedAttributes\s*\(\s*\)\s*\{\s*return\s*\[([^\]]+)\]",
            content,
            re.DOTALL
        )
        observed_attrs = []
        if observed_match:
            attrs_str = observed_match.group(1)
            # Extract quoted strings
            observed_attrs = re.findall(r"['\"]([^'\"]+)['\"]", attrs_str)

        # Parse JSDoc to extract attribute info
        # Look for @param or attribute descriptions in JSDoc
        jsdoc_match = re.search(r'/\*\*([^*]|\*(?!/))*\*/', content, re.DOTALL)
        if jsdoc_match:
            jsdoc = jsdoc_match.group(0)

            # Extract main description (first line after /**)
            desc_match = re.search(r'/\*\*\s*\n?\s*\*?\s*([^\n@*]+)', jsdoc)
            if desc_match:
                description = desc_match.group(1).strip()

            # Extract attributes section
            # Attributes:
            #   - variant: primary | secondary | ...
            #   - size: sm | default | lg
            attrs_section = re.search(
                r'Attributes:\s*\n((?:\s*\*?\s*-[^\n]+\n?)+)',
                jsdoc,
                re.IGNORECASE
            )

            if attrs_section:
                attr_lines = attrs_section.group(1)
                for line in attr_lines.split('\n'):
                    line = line.strip().lstrip('*').strip()
                    if not line.startswith('-'):
                        continue

                    # Parse: - attr_name: value1 | value2 | ...
                    attr_match = re.match(r'-\s*(\w+):\s*(.+)', line)
                    if attr_match:
                        attr_name = attr_match.group(1)
                        attr_values = attr_match.group(2).strip()

                        # Parse options from pipe-separated values
                        options = []
                        if '|' in attr_values:
                            options = [v.strip() for v in attr_values.split('|')]

                        prop_type = 'string'
                        if attr_values == 'boolean':
                            prop_type = 'boolean'
                        elif options:
                            prop_type = ' | '.join(f"'{o}'" for o in options)

                        props.append(PropInfo(
                            name=attr_name,
                            type=prop_type,
                            required=False,
                            options=options,
                        ))

        # Add any observed attributes not found in JSDoc
        jsdoc_attr_names = {p.name for p in props}
        for attr in observed_attrs:
            if attr not in jsdoc_attr_names:
                # Skip ARIA attributes for story generation
                if attr.startswith('aria-') or attr == 'tabindex':
                    continue
                props.append(PropInfo(
                    name=attr,
                    type='string',
                    required=False,
                ))

        # Check for slot usage (Web Component children)
        has_children = '<slot' in content or '<slot>' in content

        return props, tag_name, description, has_children

    def _parse_react_props(self, content: str) -> List[PropInfo]:
        """Parse TypeScript interface/type for React component props."""
        props = []

        # Extract props from interface/type
        # interface ButtonProps { variant?: 'primary' | 'secondary'; ... }
        props_pattern = re.compile(
            r'(?:interface|type)\s+\w*Props\s*(?:=\s*)?\{([^}]+)\}',
            re.DOTALL
        )

        props_match = props_pattern.search(content)
        if props_match:
            props_content = props_match.group(1)

            # Parse each prop line
            for line in props_content.split('\n'):
                line = line.strip()
                if not line or line.startswith('//'):
                    continue

                # Match: propName?: type; or propName: type;
                prop_match = re.match(
                    r'(\w+)(\?)?:\s*([^;/]+)',
                    line
                )
                if prop_match:
                    prop_name = prop_match.group(1)
                    is_optional = prop_match.group(2) == '?'
                    prop_type = prop_match.group(3).strip()

                    # Extract options from union types
                    options = []
                    if '|' in prop_type:
                        # 'primary' | 'secondary' | 'ghost'
                        options = [
                            o.strip().strip("'\"")
                            for o in prop_type.split('|')
                            if o.strip().startswith(("'", '"'))
                        ]

                    props.append(PropInfo(
                        name=prop_name,
                        type=prop_type,
                        required=not is_optional,
                        options=options,
                    ))

        return props

    @staticmethod
    def _tag_to_pascal(tag_name: str) -> str:
        """Convert kebab-case tag name to PascalCase (e.g., 'ds-button' -> 'DsButton')."""
        return ''.join(word.capitalize() for word in tag_name.split('-'))

    def _generate_csf3(self, meta: ComponentMeta, include_variants: bool) -> str:
        """Generate CSF3 format story."""
        if meta.component_type == ComponentType.WEB_COMPONENT:
            return self._generate_csf3_web_component(meta, include_variants)
        return self._generate_csf3_react(meta, include_variants)

    def _generate_csf3_react(self, meta: ComponentMeta, include_variants: bool) -> str:
        """Generate CSF3 format story for React components."""
        lines = [
            f"import type {{ Meta, StoryObj }} from '@storybook/react';",
            f"import {{ {meta.name} }} from './{meta.name}';",
            "",
            f"const meta: Meta<typeof {meta.name}> = {{",
            f"  title: 'Components/{meta.name}',",
            f"  component: {meta.name},",
            "  parameters: {",
            "    layout: 'centered',",
            "  },",
            "  tags: ['autodocs'],",
        ]

        # Add argTypes for props with options
        arg_types = []
        for prop in meta.props:
            if prop.options:
                arg_types.append(
                    f"    {prop.name}: {{\n"
                    f"      options: {prop.options},\n"
                    f"      control: {{ type: 'select' }},\n"
                    f"    }},"
                )

        if arg_types:
            lines.append("  argTypes: {")
            lines.extend(arg_types)
            lines.append("  },")

        lines.extend([
            "};",
            "",
            "export default meta;",
            f"type Story = StoryObj<typeof {meta.name}>;",
            "",
        ])

        # Generate default story
        default_args = self._get_default_args(meta)
        lines.extend([
            "export const Default: Story = {",
            "  args: {",
        ])
        for key, value in default_args.items():
            lines.append(f"    {key}: {value},")
        lines.extend([
            "  },",
            "};",
        ])

        # Generate variant stories
        if include_variants:
            variant_prop = next(
                (p for p in meta.props if p.name == 'variant' and p.options),
                None
            )
            if variant_prop:
                for variant in variant_prop.options:
                    story_name = variant.title().replace('-', '').replace('_', '')
                    lines.extend([
                        "",
                        f"export const {story_name}: Story = {{",
                        "  args: {",
                        f"    ...Default.args,",
                        f"    variant: '{variant}',",
                        "  },",
                        "};",
                    ])

            # Size variants
            size_prop = next(
                (p for p in meta.props if p.name == 'size' and p.options),
                None
            )
            if size_prop:
                for size in size_prop.options:
                    story_name = f"Size{size.title()}"
                    lines.extend([
                        "",
                        f"export const {story_name}: Story = {{",
                        "  args: {",
                        f"    ...Default.args,",
                        f"    size: '{size}',",
                        "  },",
                        "};",
                    ])

            # Disabled state
            disabled_prop = next(
                (p for p in meta.props if p.name == 'disabled'),
                None
            )
            if disabled_prop:
                lines.extend([
                    "",
                    "export const Disabled: Story = {",
                    "  args: {",
                    "    ...Default.args,",
                    "    disabled: true,",
                    "  },",
                    "};",
                ])

        return "\n".join(lines)

    def _generate_csf3_web_component(self, meta: ComponentMeta, include_variants: bool) -> str:
        """Generate CSF3 format story for Web Components using @storybook/web-components."""
        tag_name = meta.tag_name or meta.name.lower()

        lines = [
            "/**",
            f" * Storybook stories for <{tag_name}> Web Component",
            f" * Auto-generated by DSS Storybook Generator",
            " */",
            "",
            "// Import the component to ensure it's registered",
            f"import './{meta.path.split('/')[-1]}';",
            "",
        ]

        # Build argTypes object
        arg_types_entries = []
        for prop in meta.props:
            if prop.options:
                options_str = json.dumps(prop.options)
                arg_types_entries.append(
                    f"  {prop.name}: {{\n"
                    f"    options: {options_str},\n"
                    f"    control: {{ type: 'select' }},\n"
                    f"    description: '{prop.description or prop.name + ' attribute'}',\n"
                    f"  }}"
                )
            elif prop.type == 'boolean':
                arg_types_entries.append(
                    f"  {prop.name}: {{\n"
                    f"    control: {{ type: 'boolean' }},\n"
                    f"    description: '{prop.description or prop.name + ' attribute'}',\n"
                    f"  }}"
                )
            else:
                arg_types_entries.append(
                    f"  {prop.name}: {{\n"
                    f"    control: {{ type: 'text' }},\n"
                    f"    description: '{prop.description or prop.name + ' attribute'}',\n"
                    f"  }}"
                )

        lines.extend([
            "export default {",
            f"  title: 'Web Components/{meta.name}',",
            "  parameters: {",
            "    layout: 'centered',",
            "    docs: {",
            f"      description: {{ component: `{meta.description or f'{meta.name} Web Component'}` }},",
            "    },",
            "  },",
            "  tags: ['autodocs'],",
            "  argTypes: {",
        ])
        lines.append(",\n".join(arg_types_entries))
        lines.extend([
            "  },",
            "};",
            "",
        ])

        # Generate render function
        lines.extend([
            "/**",
            " * Render function that creates the Web Component with attributes",
            " */",
            "const render = (args) => {",
            f"  const el = document.createElement('{tag_name}');",
            "",
            "  // Set attributes from args",
            "  Object.entries(args).forEach(([key, value]) => {",
            "    if (key === 'children' || key === 'slot') {",
            "      el.innerHTML = value;",
            "    } else if (typeof value === 'boolean') {",
            "      if (value) el.setAttribute(key, '');",
            "      else el.removeAttribute(key);",
            "    } else if (value !== undefined && value !== null) {",
            "      el.setAttribute(key, String(value));",
            "    }",
            "  });",
            "",
            "  return el;",
            "};",
            "",
        ])

        # Get default args
        default_args = self._get_default_args_web_component(meta)
        default_args_str = json.dumps(default_args, indent=2).replace('\n', '\n  ')

        # Generate default story
        lines.extend([
            "/**",
            " * Default story showing the component in its default state",
            " */",
            "export const Default = {",
            "  render,",
            f"  args: {default_args_str},",
            "};",
        ])

        # Generate variant stories
        if include_variants:
            variant_prop = next(
                (p for p in meta.props if p.name == 'variant' and p.options),
                None
            )
            if variant_prop:
                for variant in variant_prop.options:
                    story_name = variant.title().replace('-', '').replace('_', '')
                    lines.extend([
                        "",
                        f"export const {story_name} = {{",
                        "  render,",
                        "  args: {",
                        "    ...Default.args,",
                        f"    variant: '{variant}',",
                        "  },",
                        "};",
                    ])

            # Size variants
            size_prop = next(
                (p for p in meta.props if p.name == 'size' and p.options),
                None
            )
            if size_prop:
                for size in size_prop.options:
                    story_name = f"Size{size.title().replace('-', '')}"
                    lines.extend([
                        "",
                        f"export const {story_name} = {{",
                        "  render,",
                        "  args: {",
                        "    ...Default.args,",
                        f"    size: '{size}',",
                        "  },",
                        "};",
                    ])

            # Disabled state
            disabled_prop = next(
                (p for p in meta.props if p.name == 'disabled'),
                None
            )
            if disabled_prop:
                lines.extend([
                    "",
                    "export const Disabled = {",
                    "  render,",
                    "  args: {",
                    "    ...Default.args,",
                    "    disabled: true,",
                    "  },",
                    "};",
                ])

            # Loading state (common for buttons)
            loading_prop = next(
                (p for p in meta.props if p.name == 'loading'),
                None
            )
            if loading_prop:
                lines.extend([
                    "",
                    "export const Loading = {",
                    "  render,",
                    "  args: {",
                    "    ...Default.args,",
                    "    loading: true,",
                    "  },",
                    "};",
                ])

        return "\n".join(lines)

    def _get_default_args_web_component(self, meta: ComponentMeta) -> Dict[str, Any]:
        """Get default args for a Web Component."""
        args = {}

        for prop in meta.props:
            if prop.name == 'variant' and prop.options:
                args['variant'] = prop.options[0]
            elif prop.name == 'size' and prop.options:
                args['size'] = prop.options[0]
            elif prop.name == 'type' and prop.options:
                args['type'] = prop.options[0]
            elif prop.name == 'disabled':
                args['disabled'] = False
            elif prop.name == 'loading':
                args['loading'] = False

        # Add children/slot content for components with slots
        if meta.has_children:
            args['children'] = f'Click me'

        return args

    def _generate_csf2(self, meta: ComponentMeta, include_variants: bool) -> str:
        """Generate CSF2 format story."""
        lines = [
            f"import React from 'react';",
            f"import {{ {meta.name} }} from './{meta.name}';",
            "",
            "export default {",
            f"  title: 'Components/{meta.name}',",
            f"  component: {meta.name},",
            "};",
            "",
            f"const Template = (args) => <{meta.name} {{...args}} />;",
            "",
            "export const Default = Template.bind({});",
            "Default.args = {",
        ]

        default_args = self._get_default_args(meta)
        for key, value in default_args.items():
            lines.append(f"  {key}: {value},")

        lines.append("};")

        # Generate variant stories
        if include_variants:
            variant_prop = next(
                (p for p in meta.props if p.name == 'variant' and p.options),
                None
            )
            if variant_prop:
                for variant in variant_prop.options:
                    story_name = variant.title().replace('-', '').replace('_', '')
                    lines.extend([
                        "",
                        f"export const {story_name} = Template.bind({{}});",
                        f"{story_name}.args = {{",
                        f"  ...Default.args,",
                        f"  variant: '{variant}',",
                        "};",
                    ])

        return "\n".join(lines)

    def _generate_mdx(self, meta: ComponentMeta, include_variants: bool) -> str:
        """Generate MDX format story."""
        lines = [
            f"import {{ Meta, Story, Canvas, ArgsTable }} from '@storybook/blocks';",
            f"import {{ {meta.name} }} from './{meta.name}';",
            "",
            f"<Meta title=\"Components/{meta.name}\" component={{{meta.name}}} />",
            "",
            f"# {meta.name}",
            "",
        ]

        if meta.description:
            lines.extend([meta.description, ""])

        lines.extend([
            "## Default",
            "",
            "<Canvas>",
            f"  <Story name=\"Default\">",
            f"    <{meta.name}",
        ])

        default_args = self._get_default_args(meta)
        for key, value in default_args.items():
            lines.append(f"      {key}={value}")

        lines.extend([
            f"    />",
            "  </Story>",
            "</Canvas>",
            "",
            "## Props",
            "",
            f"<ArgsTable of={{{meta.name}}} />",
        ])

        return "\n".join(lines)

    def _get_default_args(self, meta: ComponentMeta) -> Dict[str, str]:
        """Get default args for a component."""
        args = {}

        for prop in meta.props:
            if prop.name == 'children' and meta.has_children:
                args['children'] = f"'{meta.name}'"
            elif prop.name == 'variant' and prop.options:
                args['variant'] = f"'{prop.options[0]}'"
            elif prop.name == 'size' and prop.options:
                args['size'] = f"'{prop.options[0]}'"
            elif prop.name == 'disabled':
                args['disabled'] = 'false'
            elif prop.name == 'onClick':
                args['onClick'] = '() => console.log("clicked")'
            elif prop.required and prop.default_value:
                args[prop.name] = prop.default_value

        # Ensure children for button-like components
        if meta.has_children and 'children' not in args:
            args['children'] = f"'{meta.name}'"

        return args

    async def generate_stories_for_directory(
        self,
        directory: str,
        template: StoryTemplate = StoryTemplate.CSF3,
        dry_run: bool = True,
    ) -> List[Dict[str, str]]:
        """
        Generate stories for all components in a directory.

        Supports both React components (.tsx, .jsx) and Web Components (.js).

        Args:
            directory: Path to component directory
            template: Story template format
            dry_run: If True, only return what would be generated

        Returns:
            List of dicts with component path and generated story
        """
        results = []
        dir_path = self.root / directory

        if not dir_path.exists():
            return results

        # Find component files (React + Web Components)
        for pattern in ['*.tsx', '*.jsx', '*.js']:
            for comp_path in dir_path.glob(pattern):
                # Skip story files, test files, index files
                if any(x in comp_path.name.lower() for x in ['.stories.', '.test.', '.spec.', 'index.']):
                    continue

                # Check if it's a valid component file
                is_valid = False

                # React components: PascalCase naming (e.g., Button.tsx)
                if comp_path.suffix in ['.tsx', '.jsx'] and comp_path.stem[0].isupper():
                    is_valid = True

                # Web Components: check file content for class extends HTMLElement
                elif comp_path.suffix == '.js':
                    try:
                        content = comp_path.read_text(encoding='utf-8', errors='ignore')
                        if 'extends HTMLElement' in content or 'customElements.define' in content:
                            is_valid = True
                            log.debug(f"Found Web Component: {comp_path}")
                    except Exception:
                        pass

                if not is_valid:
                    continue

                try:
                    rel_path = str(comp_path.relative_to(self.root))
                    story = await self.generate_story(rel_path, template)

                    # Determine story output path (use .stories.js for Web Components)
                    if comp_path.suffix == '.js':
                        story_path = comp_path.with_name(comp_path.stem + '.stories.js')
                    else:
                        story_path = comp_path.with_suffix('.stories.tsx')

                    result = {
                        'component': rel_path,
                        'story_path': str(story_path.relative_to(self.root)),
                        'story': story,
                    }

                    if not dry_run:
                        story_path.write_text(story)
                        result['written'] = True

                    results.append(result)

                except Exception as e:
                    log.error(f"Error generating story for {comp_path}: {e}")
                    results.append({
                        'component': str(comp_path),
                        'error': str(e),
                    })

        return results