#!/usr/bin/env python3
"""
DSS MCP Server

MCP (Model Context Protocol) interface for DSS. Exposes design system
operations as tools for AI agents.

Tool Categories (32 tools):
- Status & Discovery (4): get_status, list_projects, create_project, get_project
- Token Ingestion (7): ingest_css, ingest_scss, ingest_tailwind, ingest_json, merge, export, validate
- Analysis (11): discover_project, analyze_react, find_inline_styles, find_patterns, etc.
- Storybook (5): scan, generate_story, generate_batch, theme, coverage
- Utilities (5): extract_tokens, extract_components, sync_tokens, etc.
"""

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.json_store 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 DSS server status and statistics."""
    stats = get_stats()
    figma_configured = bool(FIGMA_TOKEN)

    return json.dumps({
        "status": "ready",
        "figma": "connected" if figma_configured else "mock mode",
        "mode": "live" if figma_configured else "mock",
        "statistics": stats,
        "version": "0.8.0"
    }, indent=2)


@mcp.tool()
async def list_projects() -> str:
    """List all registered 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 for design source
    """
    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)