#!/usr/bin/env python3
"""
DSS MCP Server - Design System Server for AI Agents

Exposes design system tools via MCP protocol for Claude and other AI agents.
"""

import os
import json
import logging
import sys
from pathlib import Path
from typing import Optional

# Add parent to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))

from mcp.server.fastmcp import FastMCP
from config import config
from storage.database import Projects, Components, SyncHistory, ActivityLog, get_stats
from figma.figma_tools import FigmaToolSuite

# Import new ingestion modules
from ingest import (
    DesignToken, TokenCollection, TokenMerger, MergeStrategy,
    CSSTokenSource, SCSSTokenSource, TailwindTokenSource, JSONTokenSource
)

# Import analysis modules
from analyze import (
    ProjectScanner, ReactAnalyzer, StyleAnalyzer, DependencyGraph, QuickWinFinder,
    ProjectAnalysis, QuickWin, QuickWinType, QuickWinPriority
)

# Import Storybook modules
from storybook import (
    StorybookScanner, StoryGenerator, ThemeGenerator,
    StoryInfo, StorybookConfig, StorybookTheme, StoryTemplate
)

logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")
logger = logging.getLogger("dss-mcp")

# Initialize Figma client
FIGMA_TOKEN = os.environ.get("FIGMA_TOKEN", "")
figma = FigmaToolSuite(FIGMA_TOKEN) if FIGMA_TOKEN else None

# Create MCP server - configurable host/port for remote connections
MCP_HOST = os.environ.get("DSS_MCP_HOST", "127.0.0.1")
MCP_PORT = int(os.environ.get("DSS_MCP_PORT", "3457"))
mcp = FastMCP("dss-design-system", host=MCP_HOST, port=MCP_PORT)


@mcp.tool()
async def get_status() -> str:
    """Get the current status of the Design System Server."""
    stats = get_stats()
    figma_configured = bool(FIGMA_TOKEN)

    return json.dumps({
        "status": "online",
        "figma_configured": figma_configured,
        "figma_mode": "live" if figma_configured else "mock",
        "statistics": stats,
        "version": "1.0.0"
    }, indent=2)


@mcp.tool()
async def list_projects() -> str:
    """List all design system projects."""
    projects = Projects.list_all()
    return json.dumps([p.to_dict() for p in projects], indent=2)


@mcp.tool()
async def create_project(name: str, description: str = "", figma_file_key: str = "") -> str:
    """
    Create a new design system project.

    Args:
        name: Project name
        description: Project description
        figma_file_key: Figma file key to link
    """
    project = Projects.create(name, description, figma_file_key)
    return json.dumps({"success": True, "project": project.to_dict()}, indent=2)


@mcp.tool()
async def get_project(project_id: str) -> str:
    """
    Get details of a specific project.

    Args:
        project_id: The project ID
    """
    project = Projects.get(project_id)
    if not project:
        return json.dumps({"success": False, "error": "Project not found"})
    return json.dumps({"success": True, "project": project.to_dict()}, indent=2)


@mcp.tool()
async def extract_tokens(file_key: str, format: str = "json") -> str:
    """
    Extract design tokens from a Figma file.

    Args:
        file_key: The Figma file key (from the URL, e.g., abc123XYZ from figma.com/file/abc123XYZ/...)
        format: Output format - json, css, scss, or ts (default: json)

    Returns:
        JSON with extracted tokens and formatted output
    """
    if not figma:
        # Mock mode - return sample tokens
        mock_tokens = [
            {"name": "color-primary", "value": "#3B82F6", "type": "color", "category": "colors"},
            {"name": "color-secondary", "value": "#10B981", "type": "color", "category": "colors"},
            {"name": "spacing-sm", "value": "8px", "type": "dimension", "category": "spacing"},
            {"name": "spacing-md", "value": "16px", "type": "dimension", "category": "spacing"},
            {"name": "font-size-base", "value": "16px", "type": "dimension", "category": "typography"},
        ]
        formatted = format_tokens_output(mock_tokens, format)
        return json.dumps({
            "success": True,
            "mode": "mock",
            "tokens_count": len(mock_tokens),
            "format": format,
            "tokens": mock_tokens,
            "formatted_output": formatted
        }, indent=2)

    try:
        result = await figma.extract_variables(file_key)
        tokens = result.get("tokens", [])
        formatted = format_tokens_output(tokens, format)

        return json.dumps({
            "success": True,
            "mode": "live",
            "tokens_count": len(tokens),
            "format": format,
            "tokens": tokens[:20],  # First 20 for preview
            "formatted_output": formatted[:3000] if len(formatted) > 3000 else formatted
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def extract_components(file_key: str) -> str:
    """
    Extract component definitions from a Figma file.

    Args:
        file_key: The Figma file key (from the URL)

    Returns:
        JSON with extracted components including name, variants, and properties
    """
    if not figma:
        # Mock mode
        mock_components = [
            {"name": "Button", "key": "btn-1", "description": "Primary button component", "variants": ["primary", "secondary", "ghost"]},
            {"name": "Input", "key": "inp-1", "description": "Text input field", "variants": ["default", "error", "disabled"]},
            {"name": "Card", "key": "card-1", "description": "Content card container", "variants": ["elevated", "outlined"]},
        ]
        return json.dumps({
            "success": True,
            "mode": "mock",
            "components_count": len(mock_components),
            "components": mock_components
        }, indent=2)

    try:
        result = await figma.extract_components(file_key)
        components = result.get("components", [])

        return json.dumps({
            "success": True,
            "mode": "live",
            "components_count": len(components),
            "components": components
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def generate_component_code(
    component_name: str,
    framework: str = "react",
    tokens: str = ""
) -> str:
    """
    Generate code for a UI component.

    Args:
        component_name: Name of the component (e.g., Button, Card, Input)
        framework: Target framework - react, vue, svelte, or webcomponent (default: react)
        tokens: Optional JSON string of design tokens to use

    Returns:
        Generated component code in the specified framework
    """
    # Parse tokens if provided
    token_vars = {}
    if tokens:
        try:
            token_list = json.loads(tokens)
            for t in token_list:
                token_vars[t["name"]] = t["value"]
        except (json.JSONDecodeError, KeyError, TypeError) as e:
            # Invalid token format, continue with empty token_vars
            pass

    # Generate code based on framework
    if framework == "react":
        code = generate_react_component(component_name, token_vars)
    elif framework == "vue":
        code = generate_vue_component(component_name, token_vars)
    elif framework == "svelte":
        code = generate_svelte_component(component_name, token_vars)
    else:
        code = generate_webcomponent(component_name, token_vars)

    return json.dumps({
        "success": True,
        "component": component_name,
        "framework": framework,
        "code": code
    }, indent=2)


@mcp.tool()
async def sync_tokens_to_file(
    file_key: str,
    output_path: str,
    format: str = "css"
) -> str:
    """
    Extract tokens from Figma and save to a file.

    Args:
        file_key: The Figma file key
        output_path: Path to save the tokens file (relative or absolute)
        format: Output format - css, scss, json, or ts (default: css)

    Returns:
        Confirmation of sync with token count
    """
    # First extract tokens
    result = json.loads(await extract_tokens(file_key, format))

    if not result.get("success"):
        return json.dumps(result)

    try:
        formatted = result.get("formatted_output", "")

        # Ensure directory exists
        path = Path(output_path)
        path.parent.mkdir(parents=True, exist_ok=True)

        with open(path, "w") as f:
            f.write(formatted)

        # Log to sync history
        SyncHistory.log(
            project_id="cli",
            source="figma",
            target=str(path),
            tokens_count=result.get("tokens_count", 0),
            status="success"
        )

        return json.dumps({
            "success": True,
            "tokens_synced": result.get("tokens_count", 0),
            "output_file": str(path),
            "format": format
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def get_sync_history(limit: int = 10) -> str:
    """
    Get recent token sync history.

    Args:
        limit: Maximum number of records to return (default: 10)

    Returns:
        List of recent sync operations
    """
    history = SyncHistory.list(limit=limit)
    return json.dumps([h.to_dict() for h in history], indent=2)


@mcp.tool()
async def get_activity(limit: int = 20) -> str:
    """
    Get recent activity log.

    Args:
        limit: Maximum number of activities to return (default: 20)

    Returns:
        List of recent activities
    """
    activities = ActivityLog.list(limit=limit)
    return json.dumps([a.to_dict() for a in activities], indent=2)


# === NEW: Multi-Source Token Ingestion Tools ===

@mcp.tool()
async def ingest_css_tokens(source: str) -> str:
    """
    Extract design tokens from CSS custom properties.

    Parses CSS files for :root { --var-name: value; } declarations
    and converts them to standardized design tokens.

    Args:
        source: File path to CSS file or CSS content string

    Returns:
        JSON with extracted tokens including name, value, type, and source attribution
    """
    try:
        css_source = CSSTokenSource()
        collection = await css_source.extract(source)

        return json.dumps({
            "success": True,
            "source_type": "css",
            "tokens_count": len(collection),
            "tokens": [
                {
                    "name": t.name,
                    "value": t.value,
                    "type": t.type.value,
                    "category": t.category.value,
                    "source": t.source,
                }
                for t in collection.tokens
            ],
            "summary": collection.summary(),
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def ingest_scss_tokens(source: str) -> str:
    """
    Extract design tokens from SCSS/Sass variables.

    Parses SCSS files for $variable: value; declarations
    including map variables like $colors: (primary: #3B82F6);

    Args:
        source: File path to SCSS file or SCSS content string

    Returns:
        JSON with extracted tokens
    """
    try:
        scss_source = SCSSTokenSource()
        collection = await scss_source.extract(source)

        return json.dumps({
            "success": True,
            "source_type": "scss",
            "tokens_count": len(collection),
            "tokens": [
                {
                    "name": t.name,
                    "value": t.value,
                    "type": t.type.value,
                    "category": t.category.value,
                    "source": t.source,
                }
                for t in collection.tokens
            ],
            "summary": collection.summary(),
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def ingest_tailwind_tokens(source: str) -> str:
    """
    Extract design tokens from Tailwind CSS configuration.

    Parses tailwind.config.js/ts files for theme values
    including colors, spacing, typography, etc.

    Args:
        source: Path to tailwind.config.js or directory containing it

    Returns:
        JSON with extracted tokens
    """
    try:
        tailwind_source = TailwindTokenSource()
        collection = await tailwind_source.extract(source)

        return json.dumps({
            "success": True,
            "source_type": "tailwind",
            "tokens_count": len(collection),
            "tokens": [
                {
                    "name": t.name,
                    "value": t.value,
                    "type": t.type.value,
                    "category": t.category.value,
                    "source": t.source,
                }
                for t in collection.tokens
            ],
            "summary": collection.summary(),
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def ingest_json_tokens(source: str) -> str:
    """
    Extract design tokens from JSON token files.

    Supports multiple formats:
    - W3C Design Tokens format ($value, $type)
    - Style Dictionary format (value, comment)
    - Tokens Studio format
    - Generic nested JSON

    Args:
        source: Path to JSON file or JSON content string

    Returns:
        JSON with extracted tokens
    """
    try:
        json_source = JSONTokenSource()
        collection = await json_source.extract(source)

        return json.dumps({
            "success": True,
            "source_type": "json",
            "tokens_count": len(collection),
            "tokens": [
                {
                    "name": t.name,
                    "value": t.value,
                    "type": t.type.value,
                    "category": t.category.value,
                    "source": t.source,
                }
                for t in collection.tokens
            ],
            "summary": collection.summary(),
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def merge_tokens(
    sources: str,
    strategy: str = "last"
) -> str:
    """
    Merge tokens from multiple sources with conflict resolution.

    Combines tokens from different sources (Figma, CSS, SCSS, etc.)
    into a unified token collection.

    Args:
        sources: JSON array of source configs, e.g.:
                 [{"type": "css", "path": "/path/to/tokens.css"},
                  {"type": "scss", "path": "/path/to/variables.scss"}]
        strategy: Conflict resolution strategy:
                 - "first": Keep first occurrence
                 - "last": Keep last occurrence (default)
                 - "prefer_figma": Prefer Figma sources
                 - "prefer_code": Prefer CSS/SCSS sources
                 - "merge_metadata": Merge metadata, keep latest value

    Returns:
        JSON with merged tokens and conflict report
    """
    try:
        source_configs = json.loads(sources)
        collections = []

        # Extract tokens from each source
        for config in source_configs:
            source_type = config.get("type", "").lower()
            source_path = config.get("path", "")

            if source_type == "css":
                extractor = CSSTokenSource()
            elif source_type == "scss":
                extractor = SCSSTokenSource()
            elif source_type == "tailwind":
                extractor = TailwindTokenSource()
            elif source_type == "json":
                extractor = JSONTokenSource()
            else:
                continue

            collection = await extractor.extract(source_path)
            collections.append(collection)

        if not collections:
            return json.dumps({
                "success": False,
                "error": "No valid sources provided"
            })

        # Merge collections
        merge_strategy = MergeStrategy(strategy)
        merger = TokenMerger(strategy=merge_strategy)
        result = merger.merge(collections)

        return json.dumps({
            "success": True,
            "tokens_count": len(result.collection),
            "tokens": [
                {
                    "name": t.name,
                    "value": t.value,
                    "type": t.type.value,
                    "category": t.category.value,
                    "source": t.source,
                }
                for t in result.collection.tokens
            ],
            "conflicts": [
                {
                    "token_name": c.token_name,
                    "existing_source": c.existing.source,
                    "incoming_source": c.incoming.source,
                    "resolution": c.resolution,
                }
                for c in result.conflicts
            ],
            "stats": result.stats,
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def export_tokens(
    source: str,
    format: str = "css",
    output_path: str = ""
) -> str:
    """
    Export tokens to various formats.

    Takes tokens from any source and exports to CSS, SCSS, JSON, TypeScript, or Tailwind.

    Args:
        source: Path to token source file (CSS, SCSS, JSON, etc.) or JSON token array
        format: Output format: css, scss, json, ts, tailwind, w3c
        output_path: Optional path to save output file

    Returns:
        Formatted tokens as string, optionally saved to file
    """
    try:
        # Parse source tokens
        if source.startswith('[') or source.startswith('{'):
            # JSON content
            json_source = JSONTokenSource()
            collection = await json_source.extract(source)
        elif source.endswith('.css'):
            css_source = CSSTokenSource()
            collection = await css_source.extract(source)
        elif source.endswith('.scss') or source.endswith('.sass'):
            scss_source = SCSSTokenSource()
            collection = await scss_source.extract(source)
        else:
            json_source = JSONTokenSource()
            collection = await json_source.extract(source)

        # Format output
        if format == "css":
            output = collection.to_css()
        elif format == "scss":
            output = collection.to_scss()
        elif format == "ts":
            output = collection.to_typescript()
        elif format == "tailwind":
            output = collection.to_tailwind_config()
        elif format == "w3c":
            output = collection.to_json()
        else:  # json
            output = json.dumps([
                {"name": t.name, "value": t.value, "type": t.type.value}
                for t in collection.tokens
            ], indent=2)

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

        return json.dumps({
            "success": True,
            "format": format,
            "tokens_count": len(collection),
            "output": output[:5000] if len(output) > 5000 else output,
            "output_path": output_path if output_path else None,
            "truncated": len(output) > 5000,
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def validate_tokens(source: str) -> str:
    """
    Validate a token collection for issues.

    Checks for:
    - Duplicate token names
    - Invalid values
    - Missing required properties
    - Naming convention issues
    - Deprecated tokens

    Args:
        source: Path to token file or JSON token content

    Returns:
        Validation report with issues and warnings
    """
    try:
        # Parse tokens
        if source.endswith('.css'):
            extractor = CSSTokenSource()
        elif source.endswith('.scss'):
            extractor = SCSSTokenSource()
        else:
            extractor = JSONTokenSource()

        collection = await extractor.extract(source)

        issues = []
        warnings = []

        # Check for duplicates
        duplicates = collection.get_duplicates()
        for name, tokens in duplicates.items():
            issues.append({
                "type": "duplicate",
                "token": name,
                "count": len(tokens),
                "sources": [t.source for t in tokens],
            })

        # Check for naming issues
        for token in collection.tokens:
            # Check for spaces
            if ' ' in token.name:
                issues.append({
                    "type": "invalid_name",
                    "token": token.name,
                    "message": "Token name contains spaces",
                })

            # Check for uppercase
            if token.name != token.name.lower():
                warnings.append({
                    "type": "naming_convention",
                    "token": token.name,
                    "message": "Token name should be lowercase",
                })

            # Check for deprecated
            if token.deprecated:
                warnings.append({
                    "type": "deprecated",
                    "token": token.name,
                    "message": token.deprecated_message or "Token is deprecated",
                })

        return json.dumps({
            "success": True,
            "valid": len(issues) == 0,
            "tokens_count": len(collection),
            "issues": issues,
            "warnings": warnings,
            "summary": {
                "issues_count": len(issues),
                "warnings_count": len(warnings),
                "categories": list(cat.value for cat in collection.get_categories()),
            }
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


# === NEW: Code Analysis & Intelligence Tools ===

@mcp.tool()
async def discover_project(path: str) -> str:
    """
    Analyze a project to discover its structure, framework, and styling approach.

    Scans the codebase to identify:
    - Framework (React, Next.js, Vue, etc.)
    - Styling approaches (CSS modules, styled-components, Tailwind, etc.)
    - Component count and locations
    - Style file inventory

    Args:
        path: Path to the project root directory

    Returns:
        JSON with complete project analysis including framework, styling, and file inventory
    """
    try:
        scanner = ProjectScanner(path)
        analysis = await scanner.scan()

        return json.dumps({
            "success": True,
            "project_path": analysis.project_path,
            "framework": analysis.framework.value,
            "framework_version": analysis.framework_version,
            "styling_approaches": [sp.to_dict() for sp in analysis.styling_approaches],
            "primary_styling": analysis.primary_styling.value if analysis.primary_styling else None,
            "component_count": len([f for f in analysis.style_files if f.type == 'component']),
            "style_files": [sf.to_dict() for sf in analysis.style_files[:20]],
            "stats": analysis.stats,
            "summary": analysis.summary(),
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def analyze_react_components(path: str) -> str:
    """
    Analyze React components in a project.

    Finds all React components and extracts:
    - Component names and types (functional, class, memo, forwardRef)
    - Props definitions
    - Style file dependencies
    - Inline style usage
    - Child component relationships

    Args:
        path: Path to the project or specific directory to analyze

    Returns:
        JSON with component inventory and relationships
    """
    try:
        analyzer = ReactAnalyzer(path)
        components = await analyzer.analyze()

        return json.dumps({
            "success": True,
            "components_count": len(components),
            "components": [c.to_dict() for c in components[:30]],
            "summary": {
                "total": len(components),
                "with_styles": len([c for c in components if c.has_styles]),
                "with_inline_styles": len([c for c in components if c.inline_style_count > 0]),
                "by_type": {
                    "functional": len([c for c in components if c.type == "functional"]),
                    "class": len([c for c in components if c.type == "class"]),
                    "memo": len([c for c in components if c.type == "memo"]),
                    "forwardRef": len([c for c in components if c.type == "forwardRef"]),
                }
            }
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def find_inline_styles(path: str) -> str:
    """
    Find all inline style usage in a React project.

    Scans JSX/TSX files for style={{ ... }} patterns and style={variable} usage.

    Args:
        path: Path to the project or directory to scan

    Returns:
        JSON with inline style locations and content
    """
    try:
        analyzer = ReactAnalyzer(path)
        inline_styles = await analyzer.find_inline_styles()

        # Group by file
        by_file = {}
        for style in inline_styles:
            file_path = style['file']
            if file_path not in by_file:
                by_file[file_path] = []
            by_file[file_path].append(style)

        return json.dumps({
            "success": True,
            "total_count": len(inline_styles),
            "files_affected": len(by_file),
            "inline_styles": inline_styles[:50],
            "by_file": {k: len(v) for k, v in list(by_file.items())[:20]},
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def find_style_patterns(path: str) -> str:
    """
    Identify all styling patterns used in a project.

    Detects usage of:
    - CSS Modules
    - styled-components
    - Emotion
    - Tailwind CSS
    - Inline styles
    - Regular CSS classes

    Args:
        path: Path to the project directory

    Returns:
        JSON with styling pattern inventory by type
    """
    try:
        analyzer = ReactAnalyzer(path)
        patterns = await analyzer.find_style_patterns()

        summary = {
            pattern_type: len(occurrences)
            for pattern_type, occurrences in patterns.items()
        }

        return json.dumps({
            "success": True,
            "patterns": patterns,
            "summary": summary,
            "primary": max(summary, key=summary.get) if summary else None,
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def analyze_style_values(
    path: str,
    include_inline: bool = True
) -> str:
    """
    Analyze style values across a project to find duplicates and token candidates.

    Scans CSS files and inline styles to identify:
    - Repeated color values
    - Duplicate spacing values
    - Font values that should be tokens
    - Values that appear in multiple files

    Args:
        path: Path to the project directory
        include_inline: Whether to include inline styles in analysis

    Returns:
        JSON with duplicate values and token suggestions
    """
    try:
        analyzer = StyleAnalyzer(path)
        analysis = await analyzer.analyze(include_inline=include_inline)

        return json.dumps({
            "success": True,
            "total_values": analysis['total_values_found'],
            "unique_colors": analysis['unique_colors'],
            "unique_spacing": analysis['unique_spacing'],
            "duplicates": analysis['duplicates'][:20],
            "token_candidates": [
                {
                    "value": c.value,
                    "suggested_name": c.suggested_name,
                    "category": c.category,
                    "occurrences": c.occurrences,
                    "confidence": c.confidence,
                }
                for c in analysis['token_candidates'][:15]
            ],
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def find_unused_styles(path: str) -> str:
    """
    Find CSS classes that are defined but not used in the codebase.

    Compares CSS class definitions against className usage in JS/JSX/TS/TSX files.

    Args:
        path: Path to the project directory

    Returns:
        JSON with potentially unused CSS classes
    """
    try:
        analyzer = StyleAnalyzer(path)
        unused = await analyzer.find_unused_styles()

        return json.dumps({
            "success": True,
            "count": len(unused),
            "unused_classes": unused[:30],
            "note": "Review carefully - some classes may be dynamically generated"
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def build_source_graph(path: str, depth: int = 3) -> str:
    """
    Build a dependency graph of components and styles.

    Creates a graph showing:
    - Component import relationships
    - Style file dependencies
    - Component usage (which components render other components)

    Args:
        path: Path to the project directory
        depth: Maximum depth for traversing dependencies (default: 3)

    Returns:
        JSON with nodes (files) and edges (dependencies)
    """
    try:
        graph = DependencyGraph(path)
        result = await graph.build(depth=depth)

        return json.dumps({
            "success": True,
            "stats": result['stats'],
            "nodes": result['nodes'][:50],  # Limit for response size
            "edges": result['edges'][:100],
            "component_tree": graph.get_component_tree(),
            "hubs": graph.find_hubs(min_connections=3)[:10],
            "orphans": graph.find_orphans()[:10],
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def get_quick_wins(path: str) -> str:
    """
    Identify quick improvement opportunities in a codebase.

    Finds easy wins like:
    - Inline styles that can be extracted
    - Duplicate values that should be tokens
    - Unused CSS
    - Naming inconsistencies
    - Accessibility issues

    Args:
        path: Path to the project directory

    Returns:
        JSON with prioritized list of improvement opportunities
    """
    try:
        finder = QuickWinFinder(path)
        summary = await finder.get_summary()

        return json.dumps({
            "success": True,
            "total_quick_wins": summary['total'],
            "auto_fixable": summary['auto_fixable'],
            "by_priority": summary['by_priority'],
            "by_type": summary['by_type'],
            "top_wins": summary['top_wins'],
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def get_quick_wins_report(path: str) -> str:
    """
    Generate a human-readable report of quick improvement opportunities.

    Creates a formatted report with:
    - Prioritized list of improvements
    - Affected files
    - Suggested fixes
    - Impact estimates

    Args:
        path: Path to the project directory

    Returns:
        Formatted text report of quick-wins
    """
    try:
        finder = QuickWinFinder(path)
        report = await finder.get_actionable_report()

        return json.dumps({
            "success": True,
            "report": report,
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def check_naming_consistency(path: str) -> str:
    """
    Check CSS naming convention consistency across a project.

    Identifies the primary naming convention (BEM, kebab-case, camelCase, etc.)
    and flags classes that don't follow it.

    Args:
        path: Path to the project directory

    Returns:
        JSON with naming pattern analysis and inconsistencies
    """
    try:
        analyzer = StyleAnalyzer(path)
        naming = await analyzer.analyze_naming_consistency()

        return json.dumps({
            "success": True,
            "pattern_counts": naming['pattern_counts'],
            "primary_pattern": naming['primary_pattern'],
            "inconsistencies_count": len(naming['inconsistencies']),
            "inconsistencies": naming['inconsistencies'][:20],
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


# === NEW: Storybook Integration Tools ===

@mcp.tool()
async def scan_storybook(path: str) -> str:
    """
    Scan a project for Storybook configuration and stories.

    Discovers:
    - Storybook version and configuration
    - All story files and their components
    - Story coverage statistics

    Args:
        path: Path to the project directory

    Returns:
        JSON with Storybook configuration and story inventory
    """
    try:
        scanner = StorybookScanner(path)
        result = await scanner.scan()

        return json.dumps({
            "success": True,
            "config": result.get("config"),
            "stories_count": result.get("stories_count", 0),
            "components_with_stories": result.get("components_with_stories", 0),
            "stories": result.get("stories", [])[:30],
            "by_component": {
                k: v for k, v in list(result.get("by_component", {}).items())[:20]
            },
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def generate_story(
    component_path: str,
    template: str = "csf3",
    include_variants: bool = True,
    output_path: str = "",
) -> str:
    """
    Generate a Storybook story for a React component.

    Analyzes the component's props and generates appropriate stories
    with variants for different prop combinations.

    Args:
        component_path: Path to the component file
        template: Story format - 'csf3' (default), 'csf2', or 'mdx'
        include_variants: Generate stories for prop variants (default: True)
        output_path: Optional path to write the story file

    Returns:
        Generated story code
    """
    try:
        # Determine root from component path
        comp_path = Path(component_path)
        if comp_path.is_absolute():
            root = comp_path.parent
            rel_path = comp_path.name
        else:
            root = Path.cwd()
            rel_path = component_path

        generator = StoryGenerator(str(root))
        story_template = StoryTemplate(template)

        story = await generator.generate_story(
            rel_path,
            template=story_template,
            include_variants=include_variants,
            output_path=output_path if output_path else None,
        )

        return json.dumps({
            "success": True,
            "component": component_path,
            "template": template,
            "story": story,
            "output_path": output_path if output_path else None,
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def generate_stories_batch(
    directory: str,
    template: str = "csf3",
    dry_run: bool = True,
) -> str:
    """
    Generate Storybook stories for all components in a directory.

    Scans the directory for React components and generates stories
    for each one that doesn't already have a story file.

    Args:
        directory: Path to component directory
        template: Story format - 'csf3', 'csf2', or 'mdx'
        dry_run: If True, only preview what would be generated (default: True)

    Returns:
        JSON with generated stories and their paths
    """
    try:
        dir_path = Path(directory)
        if dir_path.is_absolute():
            root = dir_path.parent
            rel_dir = dir_path.name
        else:
            root = Path.cwd()
            rel_dir = directory

        generator = StoryGenerator(str(root))
        story_template = StoryTemplate(template)

        results = await generator.generate_stories_for_directory(
            rel_dir,
            template=story_template,
            dry_run=dry_run,
        )

        return json.dumps({
            "success": True,
            "directory": directory,
            "dry_run": dry_run,
            "generated_count": len([r for r in results if "story" in r]),
            "results": results[:20],
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def generate_storybook_theme(
    source: str,
    brand_title: str = "Design System",
    base: str = "light",
    output_dir: str = "",
) -> str:
    """
    Generate Storybook theme configuration from design tokens.

    Creates theme files that style Storybook UI to match your design system.

    Args:
        source: Path to token file (CSS, SCSS, JSON) or JSON token array
        brand_title: Brand title shown in Storybook (default: "Design System")
        base: Base theme - 'light' or 'dark' (default: 'light')
        output_dir: Optional directory to write theme files

    Returns:
        Generated theme configuration files
    """
    try:
        # Parse tokens from source
        if source.startswith('[') or source.startswith('{'):
            tokens_data = json.loads(source)
            if isinstance(tokens_data, dict):
                tokens = [{"name": k, "value": v} for k, v in tokens_data.items()]
            else:
                tokens = tokens_data
        elif source.endswith('.css'):
            css_source = CSSTokenSource()
            collection = await css_source.extract(source)
            tokens = [{"name": t.name, "value": t.value} for t in collection.tokens]
        elif source.endswith('.scss'):
            scss_source = SCSSTokenSource()
            collection = await scss_source.extract(source)
            tokens = [{"name": t.name, "value": t.value} for t in collection.tokens]
        else:
            json_source = JSONTokenSource()
            collection = await json_source.extract(source)
            tokens = [{"name": t.name, "value": t.value} for t in collection.tokens]

        generator = ThemeGenerator()
        files = generator.generate_full_config(
            tokens,
            brand_title=brand_title,
            output_dir=output_dir if output_dir else None,
        )

        return json.dumps({
            "success": True,
            "tokens_used": len(tokens),
            "brand_title": brand_title,
            "base": base,
            "files": {
                name: content[:2000] if len(content) > 2000 else content
                for name, content in files.items()
            },
            "output_dir": output_dir if output_dir else None,
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool()
async def get_story_coverage(path: str) -> str:
    """
    Get Storybook story coverage statistics for a project.

    Shows which components have stories and coverage metrics.

    Args:
        path: Path to the project directory

    Returns:
        JSON with coverage statistics
    """
    try:
        scanner = StorybookScanner(path)
        coverage = await scanner.get_story_coverage()

        return json.dumps({
            "success": True,
            "total_stories": coverage.get("total_stories", 0),
            "components_covered": coverage.get("components_covered", 0),
            "average_stories_per_component": coverage.get("average_stories_per_component", 0),
            "stories_per_component": coverage.get("stories_per_component", {}),
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


# === Helper Functions ===

def format_tokens_output(tokens: list, format: str) -> str:
    """Format tokens into the specified output format."""
    if format == "css":
        lines = [":root {"]
        for t in tokens:
            name = t["name"].replace(".", "-").replace("/", "-")
            lines.append(f"  --{name}: {t['value']};")
        lines.append("}")
        return "\n".join(lines)

    elif format == "scss":
        lines = []
        for t in tokens:
            name = t["name"].replace(".", "-").replace("/", "-")
            lines.append(f"${name}: {t['value']};")
        return "\n".join(lines)

    elif format == "ts":
        lines = ["export const tokens = {"]
        for t in tokens:
            name = t["name"].replace(".", "_").replace("/", "_").replace("-", "_")
            lines.append(f'  {name}: "{t["value"]}",')
        lines.append("} as const;")
        return "\n".join(lines)

    else:  # json
        return json.dumps(tokens, indent=2)


def generate_react_component(name: str, tokens: dict) -> str:
    """Generate a React component."""
    return f'''import React from 'react';
import styles from './{name}.module.css';

interface {name}Props {{
  children?: React.ReactNode;
  variant?: 'primary' | 'secondary';
  onClick?: () => void;
}}

export const {name}: React.FC<{name}Props> = ({{
  children,
  variant = 'primary',
  onClick
}}) => {{
  return (
    <div
      className={{`${{styles.{name.lower()}}} ${{styles[variant]}}`}}
      onClick={{onClick}}
    >
      {{children}}
    </div>
  );
}};

export default {name};
'''


def generate_vue_component(name: str, tokens: dict) -> str:
    """Generate a Vue component."""
    return f'''<template>
  <div :class="['{name.lower()}', variant]" @click="$emit('click')">
    <slot />
  </div>
</template>

<script setup lang="ts">
defineProps<{{
  variant?: 'primary' | 'secondary'
}}>()

defineEmits<{{
  click: []
}}>()
</script>

<style scoped>
.{name.lower()} {{
  /* Component styles */
}}
</style>
'''


def generate_svelte_component(name: str, tokens: dict) -> str:
    """Generate a Svelte component."""
    return f'''<script lang="ts">
  export let variant: 'primary' | 'secondary' = 'primary';
</script>

<div class="{name.lower()} {{variant}}" on:click>
  <slot />
</div>

<style>
  .{name.lower()} {{
    /* Component styles */
  }}
</style>
'''


def generate_webcomponent(name: str, tokens: dict) -> str:
    """Generate a Web Component."""
    tag_name = name.lower().replace("_", "-")
    return f'''class {name} extends HTMLElement {{
  constructor() {{
    super();
    this.attachShadow({{ mode: 'open' }});
  }}

  connectedCallback() {{
    this.render();
  }}

  render() {{
    this.shadowRoot.innerHTML = `
      <style>
        :host {{
          display: block;
        }}
      </style>
      <div class="{name.lower()}">
        <slot></slot>
      </div>
    `;
  }}
}}

customElements.define('{tag_name}', {name});
'''


if __name__ == "__main__":
    # Default to stdio transport for Claude Code integration
    transport = sys.argv[1] if len(sys.argv) > 1 else "stdio"

    logger.info(f"Starting DSS MCP Server")
    logger.info(f"  Transport: {transport}")
    logger.info(f"  Host: {MCP_HOST}")
    logger.info(f"  Port: {MCP_PORT}")
    logger.info(f"  Figma: {'configured' if FIGMA_TOKEN else 'mock mode'}")

    # Run with specified transport
    mcp.run(transport=transport)