dss/tools/dss_mcp/handler.py

"""
Unified MCP Handler

Central handler for all MCP tool execution. Used by:
- Direct API calls (/api/mcp/tools/{name}/execute)
- Claude chat (inline tool execution)
- SSE streaming connections

This module ensures all MCP requests go through a single code path
for consistent logging, error handling, and security.
"""

import json
import asyncio
from typing import Dict, Any, List, Optional, Tuple
from datetime import datetime
from dataclasses import dataclass, asdict

import sys
from pathlib import Path

# Note: sys.path is set up by the importing module (server.py)
# Do NOT modify sys.path here as it causes relative import issues

from storage.json_store import Projects, ActivityLog
from .config import mcp_config, integration_config
from .context.project_context import get_context_manager, ProjectContext
from .tools.project_tools import PROJECT_TOOLS, ProjectTools
from .tools.analysis_tools import ANALYSIS_TOOLS, AnalysisTools
from .integrations.figma import FIGMA_TOOLS, FigmaTools
from .integrations.storybook import STORYBOOK_TOOLS, StorybookTools
from .integrations.jira import JIRA_TOOLS, JiraTools
from .integrations.confluence import CONFLUENCE_TOOLS, ConfluenceTools
from .integrations.translations import TRANSLATION_TOOLS, TranslationTools
from .integrations.base import CircuitBreakerOpen


@dataclass
class ToolResult:
    """Result of a tool execution"""
    tool_name: str
    success: bool
    result: Any
    error: Optional[str] = None
    duration_ms: int = 0
    timestamp: str = None

    def __post_init__(self):
        if not self.timestamp:
            self.timestamp = datetime.now().isoformat()

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


@dataclass
class MCPContext:
    """Context for MCP operations"""
    project_id: str
    user_id: Optional[int] = None
    session_id: Optional[str] = None


class MCPHandler:
    """
    Unified MCP tool handler.

    Provides:
    - Tool discovery (list all available tools)
    - Tool execution with proper context
    - Integration management
    - Logging and metrics
    """

    def __init__(self):
        self.context_manager = get_context_manager()
        self._tool_registry: Dict[str, Dict[str, Any]] = {}
        self._initialize_tools()

    def _initialize_tools(self):
        """Initialize tool registry with all available tools"""
        # Register base project tools
        for tool in PROJECT_TOOLS:
            self._tool_registry[tool.name] = {
                "tool": tool,
                "category": "project",
                "requires_integration": False
            }

        # Register analysis tools
        for tool in ANALYSIS_TOOLS:
            self._tool_registry[tool.name] = {
                "tool": tool,
                "category": "analysis",
                "requires_integration": False
            }

        # Register Figma tools
        for tool in FIGMA_TOOLS:
            self._tool_registry[tool.name] = {
                "tool": tool,
                "category": "figma",
                "requires_integration": True,
                "integration_type": "figma"
            }

        # Register Storybook tools
        for tool in STORYBOOK_TOOLS:
            self._tool_registry[tool.name] = {
                "tool": tool,
                "category": "storybook",
                "requires_integration": False
            }

        # Register Jira tools
        for tool in JIRA_TOOLS:
            self._tool_registry[tool.name] = {
                "tool": tool,
                "category": "jira",
                "requires_integration": True,
                "integration_type": "jira"
            }

        # Register Confluence tools
        for tool in CONFLUENCE_TOOLS:
            self._tool_registry[tool.name] = {
                "tool": tool,
                "category": "confluence",
                "requires_integration": True,
                "integration_type": "confluence"
            }

        # Register Translation tools
        for tool in TRANSLATION_TOOLS:
            self._tool_registry[tool.name] = {
                "tool": tool,
                "category": "translations",
                "requires_integration": False
            }

    def list_tools(self, include_details: bool = False) -> Dict[str, Any]:
        """
        List all available MCP tools.

        Args:
            include_details: Include full tool schemas

        Returns:
            Tool listing by category
        """
        tools_by_category = {}

        for name, info in self._tool_registry.items():
            category = info["category"]
            if category not in tools_by_category:
                tools_by_category[category] = []

            tool_info = {
                "name": name,
                "description": info["tool"].description,
                "requires_integration": info.get("requires_integration", False)
            }

            if include_details:
                tool_info["input_schema"] = info["tool"].inputSchema

            tools_by_category[category].append(tool_info)

        return {
            "tools": tools_by_category,
            "total_count": len(self._tool_registry)
        }

    def get_tool_info(self, tool_name: str) -> Optional[Dict[str, Any]]:
        """Get information about a specific tool"""
        if tool_name not in self._tool_registry:
            return None

        info = self._tool_registry[tool_name]
        return {
            "name": tool_name,
            "description": info["tool"].description,
            "category": info["category"],
            "input_schema": info["tool"].inputSchema,
            "requires_integration": info.get("requires_integration", False),
            "integration_type": info.get("integration_type")
        }

    async def execute_tool(
        self,
        tool_name: str,
        arguments: Dict[str, Any],
        context: MCPContext
    ) -> ToolResult:
        """
        Execute an MCP tool.

        Args:
            tool_name: Name of the tool to execute
            arguments: Tool arguments
            context: MCP context (project_id, user_id)

        Returns:
            ToolResult with success/failure and data
        """
        start_time = datetime.now()

        # Check if tool exists
        if tool_name not in self._tool_registry:
            return ToolResult(
                tool_name=tool_name,
                success=False,
                result=None,
                error=f"Unknown tool: {tool_name}"
            )

        tool_info = self._tool_registry[tool_name]
        category = tool_info["category"]

        try:
            # Execute based on category
            if category == "project":
                result = await self._execute_project_tool(tool_name, arguments, context)
            elif category == "analysis":
                result = await self._execute_analysis_tool(tool_name, arguments, context)
            elif category == "figma":
                result = await self._execute_figma_tool(tool_name, arguments, context)
            elif category == "storybook":
                result = await self._execute_storybook_tool(tool_name, arguments, context)
            elif category == "jira":
                result = await self._execute_jira_tool(tool_name, arguments, context)
            elif category == "confluence":
                result = await self._execute_confluence_tool(tool_name, arguments, context)
            elif category == "translations":
                result = await self._execute_translations_tool(tool_name, arguments, context)
            else:
                result = {"error": f"Unknown tool category: {category}"}

            # Check for error in result
            success = "error" not in result
            error = result.get("error") if not success else None

            # Calculate duration
            duration_ms = int((datetime.now() - start_time).total_seconds() * 1000)

            # Log execution
            await self._log_tool_usage(
                tool_name=tool_name,
                category=category,
                project_id=context.project_id,
                user_id=context.user_id,
                success=success,
                duration_ms=duration_ms,
                error=error
            )

            return ToolResult(
                tool_name=tool_name,
                success=success,
                result=result if success else None,
                error=error,
                duration_ms=duration_ms
            )

        except CircuitBreakerOpen as e:
            duration_ms = int((datetime.now() - start_time).total_seconds() * 1000)
            return ToolResult(
                tool_name=tool_name,
                success=False,
                result=None,
                error=str(e),
                duration_ms=duration_ms
            )
        except Exception as e:
            duration_ms = int((datetime.now() - start_time).total_seconds() * 1000)
            await self._log_tool_usage(
                tool_name=tool_name,
                category=category,
                project_id=context.project_id,
                user_id=context.user_id,
                success=False,
                duration_ms=duration_ms,
                error=str(e)
            )
            return ToolResult(
                tool_name=tool_name,
                success=False,
                result=None,
                error=str(e),
                duration_ms=duration_ms
            )

    async def _execute_project_tool(
        self,
        tool_name: str,
        arguments: Dict[str, Any],
        context: MCPContext
    ) -> Dict[str, Any]:
        """Execute a project tool"""
        # Ensure project_id is set
        if "project_id" not in arguments:
            arguments["project_id"] = context.project_id

        project_tools = ProjectTools(context.user_id)
        return await project_tools.execute_tool(tool_name, arguments)

    async def _execute_analysis_tool(
        self,
        tool_name: str,
        arguments: Dict[str, Any],
        context: MCPContext
    ) -> Dict[str, Any]:
        """Execute an analysis tool"""
        # Ensure project_id is set for context if needed, though project_path is explicit
        if "project_id" not in arguments:
            arguments["project_id"] = context.project_id
            
        analysis_tools = AnalysisTools(context.user_id)
        return await analysis_tools.execute_tool(tool_name, arguments)

    async def _execute_figma_tool(
        self,
        tool_name: str,
        arguments: Dict[str, Any],
        context: MCPContext
    ) -> Dict[str, Any]:
        """Execute a Figma tool"""
        # Get Figma config
        config = await self._get_integration_config("figma", context)
        if not config:
            # Try global config
            if integration_config.FIGMA_TOKEN:
                config = {"api_token": integration_config.FIGMA_TOKEN}
            else:
                return {"error": "Figma not configured. Please add Figma API token."}

        figma_tools = FigmaTools(config)
        return await figma_tools.execute_tool(tool_name, arguments)

    async def _execute_storybook_tool(
        self,
        tool_name: str,
        arguments: Dict[str, Any],
        context: MCPContext
    ) -> Dict[str, Any]:
        """Execute a Storybook tool"""
        # Ensure project_id is set
        if "project_id" not in arguments:
            arguments["project_id"] = context.project_id

        storybook_tools = StorybookTools()
        return await storybook_tools.execute_tool(tool_name, arguments)

    async def _execute_jira_tool(
        self,
        tool_name: str,
        arguments: Dict[str, Any],
        context: MCPContext
    ) -> Dict[str, Any]:
        """Execute a Jira tool"""
        config = await self._get_integration_config("jira", context)
        if not config:
            return {"error": "Jira not configured. Please configure Jira integration."}

        jira_tools = JiraTools(config)
        return await jira_tools.execute_tool(tool_name, arguments)

    async def _execute_confluence_tool(
        self,
        tool_name: str,
        arguments: Dict[str, Any],
        context: MCPContext
    ) -> Dict[str, Any]:
        """Execute a Confluence tool"""
        config = await self._get_integration_config("confluence", context)
        if not config:
            return {"error": "Confluence not configured. Please configure Confluence integration."}

        confluence_tools = ConfluenceTools(config)
        return await confluence_tools.execute_tool(tool_name, arguments)

    async def _execute_translations_tool(
        self,
        tool_name: str,
        arguments: Dict[str, Any],
        context: MCPContext
    ) -> Dict[str, Any]:
        """Execute a Translation tool"""
        # Ensure project_id is set
        if "project_id" not in arguments:
            arguments["project_id"] = context.project_id

        translation_tools = TranslationTools()
        return await translation_tools.execute_tool(tool_name, arguments)

    async def _get_integration_config(
        self,
        integration_type: str,
        context: MCPContext
    ) -> Optional[Dict[str, Any]]:
        """Get decrypted integration config for user/project"""
        if not context.user_id or not context.project_id:
            return None

        loop = asyncio.get_event_loop()

        def get_config():
            try:
                with get_connection() as conn:
                    row = conn.execute(
                        """
                        SELECT config FROM project_integrations
                        WHERE project_id = ? AND user_id = ? AND integration_type = ? AND enabled = 1
                        """,
                        (context.project_id, context.user_id, integration_type)
                    ).fetchone()

                    if not row:
                        return None

                    encrypted_config = row["config"]

                    # Decrypt
                    cipher = mcp_config.get_cipher()
                    if cipher:
                        try:
                            decrypted = cipher.decrypt(encrypted_config.encode()).decode()
                            return json.loads(decrypted)
                        except:
                            pass

                    # Try parsing as plain JSON
                    try:
                        return json.loads(encrypted_config)
                    except:
                        return None
            except:
                return None

        return await loop.run_in_executor(None, get_config)

    async def _log_tool_usage(
        self,
        tool_name: str,
        category: str,
        project_id: str,
        user_id: Optional[int],
        success: bool,
        duration_ms: int,
        error: Optional[str] = None
    ):
        """Log tool execution to database"""
        loop = asyncio.get_event_loop()

        def log():
            try:
                with get_connection() as conn:
                    conn.execute(
                        """
                        INSERT INTO mcp_tool_usage
                        (project_id, user_id, tool_name, tool_category, duration_ms, success, error_message)
                        VALUES (?, ?, ?, ?, ?, ?, ?)
                        """,
                        (project_id, user_id, tool_name, category, duration_ms, success, error)
                    )
            except:
                pass  # Don't fail on logging errors

        await loop.run_in_executor(None, log)

    async def get_project_context(
        self,
        project_id: str,
        user_id: Optional[int] = None
    ) -> Optional[ProjectContext]:
        """Get project context for Claude system prompt"""
        return await self.context_manager.get_context(project_id, user_id)

    def get_tools_for_claude(self) -> List[Dict[str, Any]]:
        """
        Get tools formatted for Claude's tool_use feature.

        Returns:
            List of tools in Anthropic's tool format
        """
        tools = []
        for name, info in self._tool_registry.items():
            tools.append({
                "name": name,
                "description": info["tool"].description,
                "input_schema": info["tool"].inputSchema
            })
        return tools


# Singleton instance
_mcp_handler: Optional[MCPHandler] = None


def get_mcp_handler() -> MCPHandler:
    """Get singleton MCP handler instance"""
    global _mcp_handler
    if _mcp_handler is None:
        _mcp_handler = MCPHandler()
    return _mcp_handler