dss/tools/storybook/generator.py

"""
Storybook Story Generator

Generates Storybook stories from React components.
"""

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


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


@dataclass
class PropInfo:
    """Information about a component prop."""
    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:
    """Metadata about a component for story generation."""
    name: str
    path: str
    props: List[PropInfo] = field(default_factory=list)
    description: str = ""
    has_children: bool = False


class StoryGenerator:
    """
    Generates Storybook stories from component information.
    """

    def __init__(self, root_path: str):
        self.root = Path(root_path).resolve()

    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 React component to extract metadata."""
        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
        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,
                    ))

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

        # Extract component description from JSDoc
        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,
        )

    def _generate_csf3(self, meta: ComponentMeta, include_variants: bool) -> str:
        """Generate CSF3 format story."""
        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_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.

        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
        for pattern in ['*.tsx', '*.jsx']:
            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

                # Skip non-component files (not PascalCase)
                if not comp_path.stem[0].isupper():
                    continue

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

                    # Determine story output path
                    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:
                    results.append({
                        'component': str(comp_path),
                        'error': str(e),
                    })

        return results