From 9f0c5a5225936f28f92a2409892400a4b08960a4 Mon Sep 17 00:00:00 2001
From: sleep <abde.arrahman.dk@gmail.com>
Date: Wed, 22 Jan 2025 17:12:23 +0100
Subject: [PATCH] Added ArmoniK cluster commands

---
 src/armonik_cli/cli.py               |  1 +
 src/armonik_cli/commands/__init__.py |  3 +-
 src/armonik_cli/commands/cluster.py  | 62 ++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 src/armonik_cli/commands/cluster.py

diff --git a/src/armonik_cli/cli.py b/src/armonik_cli/cli.py
index 2816b7c..1f6fc11 100644
--- a/src/armonik_cli/cli.py
+++ b/src/armonik_cli/cli.py
@@ -15,3 +15,4 @@ def cli() -> None:
 cli.add_command(commands.sessions)
 cli.add_command(commands.tasks)
 cli.add_command(commands.partitions)
+cli.add_command(commands.cluster)
\ No newline at end of file
diff --git a/src/armonik_cli/commands/__init__.py b/src/armonik_cli/commands/__init__.py
index d1c02ab..0af0dfc 100644
--- a/src/armonik_cli/commands/__init__.py
+++ b/src/armonik_cli/commands/__init__.py
@@ -1,6 +1,7 @@
 from .sessions import sessions
 from .tasks import tasks
 from .partitions import partitions
+from .cluster import cluster
 
 
-__all__ = ["sessions", "tasks", "partitions"]
+__all__ = ["sessions", "tasks", "partitions", "cluster"]
diff --git a/src/armonik_cli/commands/cluster.py b/src/armonik_cli/commands/cluster.py
new file mode 100644
index 0000000..81b607c
--- /dev/null
+++ b/src/armonik_cli/commands/cluster.py
@@ -0,0 +1,62 @@
+import grpc
+import rich_click as click
+
+from armonik.client.versions import ArmoniKVersions
+from armonik.client.health_checks import ArmoniKHealthChecks
+from rich.panel import Panel 
+from rich.table import Table
+from rich.text import Text
+from rich import print
+
+from armonik_cli.core import base_command
+
+
+@click.group(name="cluster")
+def cluster() -> None:
+    """Manage ArmoniK cluster."""
+    pass
+
+
+@cluster.command(name="info")
+@base_command
+def cluster_info(endpoint: str, output: str, debug: bool) -> None:
+    """Get basic information on the ArmoniK cluster (endpoint, versions)"""
+    with grpc.insecure_channel(endpoint) as channel:
+        versions_client = ArmoniKVersions(channel)
+        version_info = versions_client.list_versions()
+        
+        grid = Table.grid(padding=(0, 1))
+        grid.add_column("Label", justify="left", style="cyan", no_wrap=True)
+        grid.add_column("Value", style="green", justify="left")
+        
+        grid.add_row("Endpoint:", endpoint)
+        grid.add_row("Core Version:", version_info['core'])
+        grid.add_row("API Version:", version_info['api'])
+
+        panel = Panel(grid, title="Cluster Information", border_style="blue")
+        print(panel)
+
+@cluster.command(name="health")
+@base_command
+def cluster_health(endpoint: str, output: str, debug: bool) -> None:
+    """Get information on the health of some components of the ArmoniK cluster"""
+    with grpc.insecure_channel(endpoint) as channel:
+        health_client = ArmoniKHealthChecks(channel)
+        health_status = health_client.check_health()
+        
+        grid = Table.grid(padding=(0, 1))
+        grid.add_column("Service", style="cyan", justify="left")
+        grid.add_column("Status", style="green", justify="left")
+        grid.add_column("Message", style="yellow", justify="left")
+        
+        for service_name, info in health_status.items():
+            status_style = "green" if info["status"] else "red"
+            status_text = "✓ Healthy" if info["status"] else "✗ Unhealthy"
+            grid.add_row(
+                " ".join(word.capitalize() for word in service_name.split("_")),
+                Text(status_text, style=status_style),
+                info["message"] or "-"
+            )
+
+        panel = Panel(grid, title="Health Status", border_style="blue")
+        print(panel)
\ No newline at end of file