dss/tools/dss_mcp/integrations/storybook.py

"""
Storybook Integration for MCP

Provides Storybook tools for scanning, generating stories, creating themes, and configuration.
"""

from typing import Dict, Any, Optional, List
from pathlib import Path
from mcp import types

from .base import BaseIntegration
from ..context.project_context import get_context_manager


# Storybook MCP Tool Definitions
STORYBOOK_TOOLS = [
    types.Tool(
        name="storybook_scan",
        description="Scan project for existing Storybook configuration and stories. Returns story inventory, configuration details, and coverage statistics.",
        inputSchema={
            "type": "object",
            "properties": {
                "project_id": {
                    "type": "string",
                    "description": "Project ID"
                },
                "path": {
                    "type": "string",
                    "description": "Optional: Specific path to scan (defaults to project root)"
                }
            },
            "required": ["project_id"]
        }
    ),
    types.Tool(
        name="storybook_generate_stories",
        description="Generate Storybook stories for React components. Supports CSF3, CSF2, and MDX formats with automatic prop detection.",
        inputSchema={
            "type": "object",
            "properties": {
                "project_id": {
                    "type": "string",
                    "description": "Project ID"
                },
                "component_path": {
                    "type": "string",
                    "description": "Path to component file or directory"
                },
                "template": {
                    "type": "string",
                    "description": "Story format template",
                    "enum": ["csf3", "csf2", "mdx"],
                    "default": "csf3"
                },
                "include_variants": {
                    "type": "boolean",
                    "description": "Generate variant stories (default: true)",
                    "default": True
                },
                "dry_run": {
                    "type": "boolean",
                    "description": "Preview without writing files (default: true)",
                    "default": True
                }
            },
            "required": ["project_id", "component_path"]
        }
    ),
    types.Tool(
        name="storybook_generate_theme",
        description="Generate Storybook theme configuration from design tokens. Creates manager.ts, preview.ts, and theme files.",
        inputSchema={
            "type": "object",
            "properties": {
                "project_id": {
                    "type": "string",
                    "description": "Project ID"
                },
                "brand_title": {
                    "type": "string",
                    "description": "Brand title for Storybook UI",
                    "default": "Design System"
                },
                "base_theme": {
                    "type": "string",
                    "description": "Base theme (light or dark)",
                    "enum": ["light", "dark"],
                    "default": "light"
                },
                "output_dir": {
                    "type": "string",
                    "description": "Output directory (default: .storybook)"
                },
                "write_files": {
                    "type": "boolean",
                    "description": "Write files to disk (default: false - preview only)",
                    "default": False
                }
            },
            "required": ["project_id"]
        }
    ),
    types.Tool(
        name="storybook_get_status",
        description="Get Storybook installation and configuration status for a project.",
        inputSchema={
            "type": "object",
            "properties": {
                "project_id": {
                    "type": "string",
                    "description": "Project ID"
                }
            },
            "required": ["project_id"]
        }
    ),
    types.Tool(
        name="storybook_configure",
        description="Configure or update Storybook for a project with DSS integration.",
        inputSchema={
            "type": "object",
            "properties": {
                "project_id": {
                    "type": "string",
                    "description": "Project ID"
                },
                "action": {
                    "type": "string",
                    "description": "Configuration action",
                    "enum": ["init", "update", "add_theme"],
                    "default": "init"
                },
                "options": {
                    "type": "object",
                    "description": "Configuration options",
                    "properties": {
                        "framework": {
                            "type": "string",
                            "enum": ["react", "vue", "angular"]
                        },
                        "builder": {
                            "type": "string",
                            "enum": ["vite", "webpack5"]
                        },
                        "typescript": {
                            "type": "boolean"
                        }
                    }
                }
            },
            "required": ["project_id"]
        }
    )
]


class StorybookIntegration(BaseIntegration):
    """Storybook integration wrapper for DSS tools"""

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """
        Initialize Storybook integration.

        Args:
            config: Optional Storybook configuration
        """
        super().__init__("storybook", config or {})
        self.context_manager = get_context_manager()

    async def _get_project_path(self, project_id: str) -> Path:
        """
        Get project path from context manager.

        Args:
            project_id: Project ID

        Returns:
            Project path as Path object
        """
        context = await self.context_manager.get_context(project_id)
        if not context or not context.path:
            raise ValueError(f"Project not found: {project_id}")
        return Path(context.path)

    async def scan_storybook(self, project_id: str, path: Optional[str] = None) -> Dict[str, Any]:
        """
        Scan for Storybook config and stories.

        Args:
            project_id: Project ID
            path: Optional specific path to scan

        Returns:
            Storybook scan results
        """
        try:
            from dss.storybook.scanner import StorybookScanner

            project_path = await self._get_project_path(project_id)

            # Ensure path is within project directory for security
            if path:
                scan_path = project_path / path
                # Validate path doesn't escape project directory
                if not scan_path.resolve().is_relative_to(project_path.resolve()):
                    raise ValueError("Path must be within project directory")
            else:
                scan_path = project_path

            scanner = StorybookScanner(str(scan_path))
            result = await scanner.scan() if hasattr(scanner.scan, '__await__') else scanner.scan()
            coverage = await scanner.get_story_coverage() if hasattr(scanner.get_story_coverage, '__await__') else scanner.get_story_coverage()

            return {
                "project_id": project_id,
                "path": str(scan_path),
                "config": result.get("config") if isinstance(result, dict) else None,
                "stories_count": result.get("stories_count", 0) if isinstance(result, dict) else 0,
                "components_with_stories": result.get("components_with_stories", []) if isinstance(result, dict) else [],
                "stories": result.get("stories", []) if isinstance(result, dict) else [],
                "coverage": coverage if coverage else {}
            }
        except Exception as e:
            return {
                "error": f"Failed to scan Storybook: {str(e)}",
                "project_id": project_id
            }

    async def generate_stories(
        self,
        project_id: str,
        component_path: str,
        template: str = "csf3",
        include_variants: bool = True,
        dry_run: bool = True
    ) -> Dict[str, Any]:
        """
        Generate stories for components.

        Args:
            project_id: Project ID
            component_path: Path to component file or directory
            template: Story format (csf3, csf2, mdx)
            include_variants: Whether to generate variant stories
            dry_run: Preview without writing files

        Returns:
            Generation results
        """
        try:
            from dss.storybook.generator import StoryGenerator

            project_path = await self._get_project_path(project_id)
            generator = StoryGenerator(str(project_path))

            full_path = project_path / component_path

            # Check if path exists and is directory or file
            if not full_path.exists():
                return {
                    "error": f"Path not found: {component_path}",
                    "project_id": project_id
                }

            if full_path.is_dir():
                # Generate for directory
                func = generator.generate_stories_for_directory
                if hasattr(func, '__await__'):
                    results = await func(
                        component_path,
                        template=template.upper(),
                        dry_run=dry_run
                    )
                else:
                    results = func(
                        component_path,
                        template=template.upper(),
                        dry_run=dry_run
                    )

                return {
                    "project_id": project_id,
                    "path": component_path,
                    "generated_count": len([r for r in (results if isinstance(results, list) else []) if "story" in str(r)]),
                    "results": results if isinstance(results, list) else [],
                    "dry_run": dry_run,
                    "template": template
                }
            else:
                # Generate for single file
                func = generator.generate_story
                if hasattr(func, '__await__'):
                    story = await func(
                        component_path,
                        template=template.upper(),
                        include_variants=include_variants
                    )
                else:
                    story = func(
                        component_path,
                        template=template.upper(),
                        include_variants=include_variants
                    )

                return {
                    "project_id": project_id,
                    "component": component_path,
                    "story": story,
                    "template": template,
                    "include_variants": include_variants,
                    "dry_run": dry_run
                }

        except Exception as e:
            return {
                "error": f"Failed to generate stories: {str(e)}",
                "project_id": project_id,
                "component_path": component_path
            }

    async def generate_theme(
        self,
        project_id: str,
        brand_title: str = "Design System",
        base_theme: str = "light",
        output_dir: Optional[str] = None,
        write_files: bool = False
    ) -> Dict[str, Any]:
        """
        Generate Storybook theme from design tokens.

        Args:
            project_id: Project ID
            brand_title: Brand title for Storybook
            base_theme: Base theme (light or dark)
            output_dir: Output directory for theme files
            write_files: Write files to disk or preview only

        Returns:
            Theme generation results
        """
        try:
            from dss.storybook.theme import ThemeGenerator
            from dss.themes import get_default_light_theme, get_default_dark_theme

            # Get project tokens from context
            context = await self.context_manager.get_context(project_id)
            if not context:
                return {"error": f"Project not found: {project_id}"}

            # Convert tokens to list format for ThemeGenerator
            tokens_list = [
                {"name": name, "value": token.get("value") if isinstance(token, dict) else token}
                for name, token in (context.tokens.items() if hasattr(context, 'tokens') else {}.items())
            ]

            generator = ThemeGenerator()

            if write_files and output_dir:
                # Generate and write files
                func = generator.generate_full_config
                if hasattr(func, '__await__'):
                    files = await func(
                        tokens=tokens_list,
                        brand_title=brand_title,
                        output_dir=output_dir
                    )
                else:
                    files = func(
                        tokens=tokens_list,
                        brand_title=brand_title,
                        output_dir=output_dir
                    )

                return {
                    "project_id": project_id,
                    "files_written": list(files.keys()) if isinstance(files, dict) else [],
                    "output_dir": output_dir,
                    "brand_title": brand_title
                }
            else:
                # Preview mode - generate file contents
                try:
                    func = generator.generate_from_tokens
                    if hasattr(func, '__await__'):
                        theme = await func(tokens_list, brand_title, base_theme)
                    else:
                        theme = func(tokens_list, brand_title, base_theme)
                except Exception:
                    # Fallback to default theme
                    theme_obj = get_default_light_theme() if base_theme == "light" else get_default_dark_theme()
                    theme = {
                        "name": theme_obj.name if hasattr(theme_obj, 'name') else "Default",
                        "colors": {}
                    }

                # Generate theme file content
                theme_file = f"// Storybook theme for {brand_title}\nexport default {str(theme)};"
                manager_file = f"import addons from '@storybook/addons';\nimport theme from './dss-theme';\naddons.setConfig({{ theme }});"
                preview_file = f"import '../dss-theme';\nexport default {{ parameters: {{ actions: {{ argTypesRegex: '^on[A-Z].*' }} }} }};"

                return {
                    "project_id": project_id,
                    "preview": True,
                    "brand_title": brand_title,
                    "base_theme": base_theme,
                    "files": {
                        "dss-theme.ts": theme_file,
                        "manager.ts": manager_file,
                        "preview.ts": preview_file
                    },
                    "token_count": len(tokens_list)
                }

        except Exception as e:
            return {
                "error": f"Failed to generate theme: {str(e)}",
                "project_id": project_id
            }

    async def get_status(self, project_id: str) -> Dict[str, Any]:
        """
        Get Storybook installation and configuration status.

        Args:
            project_id: Project ID

        Returns:
            Storybook status information
        """
        try:
            from dss.storybook.config import get_storybook_status

            project_path = await self._get_project_path(project_id)

            func = get_storybook_status
            if hasattr(func, '__await__'):
                status = await func(str(project_path))
            else:
                status = func(str(project_path))

            return {
                "project_id": project_id,
                "path": str(project_path),
                **(status if isinstance(status, dict) else {})
            }

        except Exception as e:
            return {
                "error": f"Failed to get Storybook status: {str(e)}",
                "project_id": project_id,
                "installed": False
            }

    async def configure(
        self,
        project_id: str,
        action: str = "init",
        options: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        Configure or update Storybook for project.

        Args:
            project_id: Project ID
            action: Configuration action (init, update, add_theme)
            options: Configuration options

        Returns:
            Configuration results
        """
        try:
            from dss.storybook.config import write_storybook_config_file

            project_path = await self._get_project_path(project_id)
            options = options or {}

            # Map action to configuration
            config = {
                "action": action,
                "framework": options.get("framework", "react"),
                "builder": options.get("builder", "vite"),
                "typescript": options.get("typescript", True)
            }

            func = write_storybook_config_file
            if hasattr(func, '__await__'):
                result = await func(str(project_path), config)
            else:
                result = func(str(project_path), config)

            return {
                "project_id": project_id,
                "action": action,
                "success": True,
                "path": str(project_path),
                "config_path": str(project_path / ".storybook"),
                "options": config
            }

        except Exception as e:
            return {
                "error": f"Failed to configure Storybook: {str(e)}",
                "project_id": project_id,
                "action": action,
                "success": False
            }


class StorybookTools:
    """MCP tool executor for Storybook integration"""

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """
        Args:
            config: Optional Storybook configuration
        """
        self.storybook = StorybookIntegration(config)

    async def execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """
        Execute Storybook tool.

        Args:
            tool_name: Name of tool to execute
            arguments: Tool arguments

        Returns:
            Tool execution result
        """
        handlers = {
            "storybook_scan": self.storybook.scan_storybook,
            "storybook_generate_stories": self.storybook.generate_stories,
            "storybook_generate_theme": self.storybook.generate_theme,
            "storybook_get_status": self.storybook.get_status,
            "storybook_configure": self.storybook.configure
        }

        handler = handlers.get(tool_name)
        if not handler:
            return {"error": f"Unknown Storybook tool: {tool_name}"}

        try:
            # Remove internal prefixes and execute
            clean_args = {k: v for k, v in arguments.items() if not k.startswith("_")}
            result = await handler(**clean_args)
            return result
        except Exception as e:
            return {"error": f"Tool execution failed: {str(e)}", "tool": tool_name}