python-sdk/pkg/hanzo-mcp/hanzo_mcp/tools/common/stats.py

"""Comprehensive system and MCP statistics."""

from datetime import datetime
from pathlib import Path
from typing import TypedDict, Unpack, final, override

import psutil
from hanzo_tools.database.database_manager import DatabaseManager
from hanzo_tools.shell.run_background import RunBackgroundTool
from mcp.server.fastmcp import Context as MCPContext

from hanzo_mcp.tools.common.auto_timeout import auto_timeout
from hanzo_mcp.tools.common.base import BaseTool
from hanzo_mcp.tools.common.context import create_tool_context
from hanzo_mcp.tools.mcp.mcp_add import McpAddTool


class StatsParams(TypedDict, total=False):
    """Parameters for stats tool."""

    pass


@final
class StatsTool(BaseTool):
    """Tool for showing comprehensive system and MCP statistics."""

    def __init__(self, db_manager: DatabaseManager = None):
        """Initialize the stats tool.

        Args:
            db_manager: Optional database manager for DB stats
        """
        self.db_manager = db_manager

    @property
    @override
    def name(self) -> str:
        """Get the tool name."""
        return "stats"

    @property
    @override
    def description(self) -> str:
        """Get the tool description."""
        return """Show comprehensive system and Hanzo AI statistics.

Displays:
- System resources (CPU, memory, disk)
- Running processes
- Database usage
- MCP server status
- Tool usage statistics
- Warnings for high resource usage

Example:
- stats
"""

    @override
    @auto_timeout("stats")
    async def call(
        self,
        ctx: MCPContext,
        **params: Unpack[StatsParams],
    ) -> str:
        """Get comprehensive statistics.

        Args:
            ctx: MCP context
            **params: Tool parameters

        Returns:
            Comprehensive statistics
        """
        tool_ctx = create_tool_context(ctx)
        await tool_ctx.set_tool_info(self.name)

        output = []
        warnings = []

        # Header
        output.append("=== Hanzo AI System Statistics ===")
        output.append(f"Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        output.append("")

        # System Resources
        output.append("=== System Resources ===")

        # CPU
        cpu_percent = psutil.cpu_percent(interval=1)
        cpu_count = psutil.cpu_count()
        output.append(f"CPU Usage: {cpu_percent}% ({cpu_count} cores)")
        if cpu_percent > 90:
            warnings.append(f"⚠️  HIGH CPU USAGE: {cpu_percent}%")

        # Memory
        memory = psutil.virtual_memory()
        memory_used_gb = memory.used / (1024**3)
        memory_total_gb = memory.total / (1024**3)
        memory_percent = memory.percent
        output.append(
            f"Memory: {memory_used_gb:.1f}/{memory_total_gb:.1f} GB ({memory_percent}%)"
        )
        if memory_percent > 90:
            warnings.append(f"⚠️  HIGH MEMORY USAGE: {memory_percent}%")

        # Disk
        disk = psutil.disk_usage("/")
        disk_used_gb = disk.used / (1024**3)
        disk_total_gb = disk.total / (1024**3)
        disk_percent = disk.percent
        disk_free_gb = disk.free / (1024**3)
        output.append(
            f"Disk: {disk_used_gb:.1f}/{disk_total_gb:.1f} GB ({disk_percent}%)"
        )
        output.append(f"Free Space: {disk_free_gb:.1f} GB")
        if disk_percent > 90:
            warnings.append(
                f"⚠️  LOW DISK SPACE: Only {disk_free_gb:.1f} GB free ({100 - disk_percent:.1f}% remaining)"
            )

        output.append("")

        # Background Processes
        output.append("=== Background Processes ===")
        processes = RunBackgroundTool.get_processes()
        running_count = 0
        total_memory_mb = 0

        if processes:
            for proc in processes.values():
                if proc.is_running():
                    running_count += 1
                    try:
                        ps_proc = psutil.Process(proc.process.pid)
                        memory_mb = ps_proc.memory_info().rss / (1024**2)
                        total_memory_mb += memory_mb
                    except Exception:
                        pass

            output.append(f"Running Processes: {running_count}")
            output.append(f"Total Memory Usage: {total_memory_mb:.1f} MB")

            # List top processes by memory
            if running_count > 0:
                output.append("\nTop Processes:")
                proc_list = []
                for proc_id, proc in processes.items():
                    if proc.is_running():
                        try:
                            ps_proc = psutil.Process(proc.process.pid)
                            memory_mb = ps_proc.memory_info().rss / (1024**2)
                            cpu = ps_proc.cpu_percent(interval=0.1)
                            proc_list.append((proc.name, memory_mb, cpu, proc_id))
                        except Exception:
                            proc_list.append((proc.name, 0, 0, proc_id))

                proc_list.sort(key=lambda x: x[1], reverse=True)
                for name, mem, cpu, pid in proc_list[:5]:
                    output.append(f"  - {name} ({pid}): {mem:.1f} MB, {cpu:.1f}% CPU")
        else:
            output.append("No background processes running")

        output.append("")

        # Database Usage
        if self.db_manager:
            output.append("=== Database Usage ===")
            db_dir = Path.home() / ".hanzo" / "db"
            total_db_size = 0

            if db_dir.exists():
                for db_file in db_dir.rglob("*.db"):
                    size = db_file.stat().st_size
                    total_db_size += size

                output.append(
                    f"Total Database Size: {total_db_size / (1024**2):.1f} MB"
                )
                output.append(f"Active Projects: {len(self.db_manager.projects)}")

                # List largest databases
                db_sizes = []
                for db_file in db_dir.rglob("*.db"):
                    size = db_file.stat().st_size / (1024**2)
                    if size > 0.1:  # Only show DBs > 100KB
                        project = db_file.parent.parent.name
                        db_type = db_file.stem
                        db_sizes.append((project, db_type, size))

                if db_sizes:
                    db_sizes.sort(key=lambda x: x[2], reverse=True)
                    output.append("\nLargest Databases:")
                    for project, db_type, size in db_sizes[:5]:
                        output.append(f"  - {project}/{db_type}: {size:.1f} MB")
            else:
                output.append("No databases found")

        output.append("")

        # MCP Servers
        output.append("=== MCP Servers ===")
        mcp_servers = McpAddTool.get_servers()
        if mcp_servers:
            running_mcp = sum(
                1 for s in mcp_servers.values() if s.get("status") == "running"
            )
            total_mcp_tools = sum(len(s.get("tools", [])) for s in mcp_servers.values())

            output.append(f"Total Servers: {len(mcp_servers)}")
            output.append(f"Running: {running_mcp}")
            output.append(f"Total Tools Available: {total_mcp_tools}")
        else:
            output.append("No MCP servers configured")

        output.append("")

        # Hanzo AI Specifics
        output.append("=== Hanzo AI ===")

        # Log directory size
        log_dir = Path.home() / ".hanzo" / "logs"
        if log_dir.exists():
            log_size = sum(f.stat().st_size for f in log_dir.rglob("*") if f.is_file())
            log_count = len(list(log_dir.rglob("*.log")))
            output.append(f"Log Files: {log_count} ({log_size / (1024**2):.1f} MB)")

            if log_size > 100 * 1024**2:  # > 100MB
                warnings.append(
                    f"⚠️  Large log directory: {log_size / (1024**2):.1f} MB"
                )

        # Config directory
        config_dir = Path.home() / ".hanzo" / "mcp"
        if config_dir.exists():
            config_count = len(list(config_dir.rglob("*.json")))
            output.append(f"Config Files: {config_count}")

        # Tool status (if available)
        # Note: Tool usage statistics tracking can be added here
        output.append("\nTool Categories:")
        output.append("  - File Operations: grep, find_files, read, write, edit")
        output.append("  - Shell: bash, run_background, processes, pkill")
        output.append("  - Database: sql_query, graph_query, vector_search")
        output.append("  - Package Runners: uvx, npx, uvx_background, npx_background")
        output.append("  - MCP Management: mcp_add, mcp_remove, mcp_stats")

        # Warnings Section
        if warnings:
            output.append("\n=== ⚠️  WARNINGS ===")
            for warning in warnings:
                output.append(warning)
            output.append("")

        # Recommendations
        output.append("=== Recommendations ===")
        if disk_free_gb < 5:
            output.append("- Free up disk space (< 5GB remaining)")
        if memory_percent > 80:
            output.append("- Close unused applications to free memory")
        if running_count > 10:
            output.append("- Consider stopping unused background processes")
        if log_size > 50 * 1024**2:
            output.append("- Clean up old log files in ~/.hanzo/logs")

        if not any(
            [
                disk_free_gb < 5,
                memory_percent > 80,
                running_count > 10,
                log_size > 50 * 1024**2,
            ]
        ):
            output.append("✅ System resources are healthy")

        return "\n".join(output)

    def register(self, mcp_server) -> None:
        """Register this tool with the MCP server."""
        pass