import psutil
import argparse
import json
from importlib.metadata import PackageNotFoundError, version


def get_system_stats():
    # Core Logic
    stats = {
        "cpu_usage": f"{psutil.cpu_percent()}%",
        "memory_info": f"{round(psutil.virtual_memory().available / (1024**3), 2)} GB available"
    }
    return stats


def run_mcp_server():
    from mcp.server.fastmcp import FastMCP

    mcp = FastMCP("system-stats")

    @mcp.tool()
    def system_stats():
        return get_system_stats()

    mcp.run(transport="stdio")


def get_cli_version():
    try:
        return version("metal-tools")
    except PackageNotFoundError:
        return "0.1.0-dev"


def main():
    # CLI layer for both `python system_stats.py` and console script entrypoint.
    parser = argparse.ArgumentParser(description="Check Metal compute resource usage")
    parser.add_argument("--json", action="store_true", help="Output in JSON format")
    parser.add_argument("--mcp", action="store_true", help="Run as an MCP stdio server")
    parser.add_argument("--version", action="version", version=f"%(prog)s {get_cli_version()}")
    args = parser.parse_args()

    if args.mcp:
        run_mcp_server()
        return

    data = get_system_stats()
    if args.json:
        print(json.dumps(data))
    else:
        print(f"CPU: {data['cpu_usage']} | RAM: {data['memory_info']}")


if __name__ == "__main__":
    main()
