#!/usr/bin/env python3
"""
SLURM Cluster Monitoring Tool
Provides comprehensive overview of cluster usage, resource allocation, and potential issues.
"""

import subprocess
import json
import sys
from collections import defaultdict
from datetime import datetime
import argparse


class ClusterMonitor:
    def __init__(self, verbose=False):
        self.verbose = verbose
        self.timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    def run_command(self, cmd):
        """Execute a shell command and return output."""
        try:
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
            if result.returncode != 0 and self.verbose:
                print(f"Warning: Command failed: {cmd}")
                print(f"Error: {result.stderr}")
            return result.stdout
        except Exception as e:
            if self.verbose:
                print(f"Error running command: {e}")
            return ""

    def get_node_summary(self):
        """Get summary of node states."""
        cmd = "sinfo -h -N -o '%T' | sort | uniq -c"
        output = self.run_command(cmd)

        states = {}
        total_nodes = 0
        for line in output.strip().split("\n"):
            if line:
                parts = line.strip().split()
                if len(parts) == 2:
                    count, state = parts
                    states[state] = int(count)
                    total_nodes += int(count)

        return states, total_nodes

    def get_gpu_usage_by_user(self):
        """Calculate GPU usage per user."""
        # Get running jobs with user and node info
        cmd = "squeue -h -t RUNNING -o '%.30u %.8D %.6C %.200R'"
        output = self.run_command(cmd)

        user_stats = defaultdict(
            lambda: {"nodes": 0, "cpus": 0, "gpus": 0, "node_list": []}
        )

        for line in output.strip().split("\n"):
            if line:
                parts = line.strip().split()
                if len(parts) >= 4:
                    user = parts[0]
                    nodes = int(parts[1])
                    cpus = int(parts[2])
                    node_list = parts[3]

                    # Assuming 8 GPUs per node for H100 nodes
                    gpus = nodes * 8

                    user_stats[user]["nodes"] += nodes
                    user_stats[user]["cpus"] += cpus
                    user_stats[user]["gpus"] += gpus
                    user_stats[user]["node_list"].append(node_list)

        return dict(user_stats)

    def get_idle_nodes(self):
        """Get list of idle nodes."""
        cmd = "sinfo -h -t idle -N -o '%N'"
        output = self.run_command(cmd)

        idle_nodes = []
        for line in output.strip().split("\n"):
            if line:
                idle_nodes.append(line.strip())

        return idle_nodes

    def get_drained_nodes(self):
        """Get drained/down nodes with reasons."""
        cmd = "sinfo -R -o '%50E %20H %10T %N'"
        output = self.run_command(cmd)

        drained = []
        lines = output.strip().split("\n")
        for line in lines[1:]:  # Skip header
            if line and ("drain" in line.lower() or "down" in line.lower()):
                drained.append(line.strip())

        return drained

    def get_queue_summary(self):
        """Get job queue summary."""
        cmd = "squeue -h -o '%T' | sort | uniq -c"
        output = self.run_command(cmd)

        queue_stats = {}
        for line in output.strip().split("\n"):
            if line:
                parts = line.strip().split()
                if len(parts) == 2:
                    count, state = parts
                    queue_stats[state] = int(count)

        return queue_stats

    def get_running_jobs_with_commands(self):
        """Get running jobs with their launch commands, workdir, and stdout."""
        # Get basic job info
        cmd = "squeue -h -t RUNNING -o '%.10i|%.15u|%.6D|%.6C|%.10M|%.50j'"
        output = self.run_command(cmd)

        jobs = []
        for line in output.strip().split("\n"):
            if line:
                parts = line.split("|")
                if len(parts) >= 6:
                    try:
                        job_id = parts[0].strip()

                        # Start with basic info from squeue
                        job_info = {
                            "id": job_id,
                            "user": parts[1].strip(),
                            "nodes": int(parts[2].strip()),
                            "cpus": int(parts[3].strip()),
                            "time": parts[4].strip(),
                            "name": parts[5].strip() if len(parts) > 5 else "N/A",
                            "command": "N/A",
                            "workdir": "N/A",
                            "stdout": "N/A",
                        }

                        # Try to get detailed info from scontrol
                        cmd_detail = f"scontrol show job {job_id} 2>/dev/null"
                        detail_output = self.run_command(cmd_detail)

                        # Parse details if available
                        if detail_output and "JobId=" in detail_output:
                            for detail_line in detail_output.split("\n"):
                                if "Command=" in detail_line:
                                    command = detail_line.split("Command=")[1].strip()
                                    if "/" in command:
                                        job_info["command"] = command.split("/")[-1]
                                    else:
                                        job_info["command"] = command
                                elif "WorkDir=" in detail_line:
                                    workdir = (
                                        detail_line.split("WorkDir=")[1]
                                        .split()[0]
                                        .strip()
                                    )
                                    # Shorten workdir path for display
                                    if len(workdir) > 30:
                                        path_parts = workdir.split("/")
                                        if len(path_parts) > 3:
                                            job_info["workdir"] = (
                                                f".../{'/'.join(path_parts[-2:])}"
                                            )
                                        else:
                                            job_info["workdir"] = workdir
                                    else:
                                        job_info["workdir"] = workdir
                                elif "StdOut=" in detail_line:
                                    stdout = (
                                        detail_line.split("StdOut=")[1]
                                        .split()[0]
                                        .strip()
                                    )
                                    # Extract just the filename
                                    if "/" in stdout:
                                        job_info["stdout"] = stdout.split("/")[-1]
                                    else:
                                        job_info["stdout"] = stdout
                        elif self.verbose:
                            print(f"Warning: Could not get details for job {job_id}")

                        jobs.append(job_info)

                    except (ValueError, IndexError) as e:
                        if self.verbose:
                            print(f"Warning: Error parsing job entry: {e}")

        return jobs

    def get_long_running_jobs(self, days=3):
        """Get jobs running longer than specified days."""
        cmd = f"squeue -h -t RUNNING -o '%.10i %.15u %.8T %.15M %.6D %.50j' | awk '$4 ~ /-/ && int($4) >= {days}'"
        output = self.run_command(cmd)

        long_jobs = []
        for line in output.strip().split("\n"):
            if line:
                long_jobs.append(line.strip())

        return long_jobs

    def get_partition_summary(self):
        """Get detailed partition information."""
        cmd = "sinfo -o '%20P %5D %10T %5c %10m %10a'"
        output = self.run_command(cmd)
        return output

    def trim_middle(self, text, max_len=20):
        """Trim long text with ... in the middle."""
        if len(text) <= max_len:
            return text
        if max_len <= 3:
            return text[:max_len]
        # Calculate how many chars to keep on each side
        keep_chars = (max_len - 3) // 2
        left_chars = keep_chars
        right_chars = max_len - 3 - left_chars
        return f"{text[:left_chars]}...{text[-right_chars:]}"

    def print_report(self):
        """Print comprehensive cluster report."""
        print("=" * 80)
        print(f"SLURM CLUSTER MONITORING REPORT - {self.timestamp}")
        print("=" * 80)

        # Node Summary
        print("\n📊 NODE STATUS SUMMARY")
        print("-" * 40)
        states, total = self.get_node_summary()

        # Define the desired order: allocated, mixed, idle, drained (and any others)
        ordered_states = ["allocated", "mixed", "idle", "drained"]

        # Print states in the specified order
        for state in ordered_states:
            if state in states:
                count = states[state]
                percentage = (count / total * 100) if total > 0 else 0
                status_icon = "✅" if state in ["idle", "mixed", "allocated"] else "⚠️"
                print(f"{status_icon} {state:15} {count:4} nodes ({percentage:5.1f}%)")

        # Print any other states that weren't in our ordered list
        for state, count in sorted(states.items()):
            if state not in ordered_states:
                percentage = (count / total * 100) if total > 0 else 0
                status_icon = "⚠️"  # Unknown states get warning icon
                print(f"{status_icon} {state:15} {count:4} nodes ({percentage:5.1f}%)")

        print(f"\n{'Total nodes:':15} {total:4}")

        # GPU Usage by User
        print("\n👥 GPU ALLOCATION BY USER")
        print("-" * 40)
        user_stats = self.get_gpu_usage_by_user()
        if user_stats:
            # Sort by GPU count
            sorted_users = sorted(
                user_stats.items(), key=lambda x: x[1]["gpus"], reverse=True
            )
            print(f"{'User':<15} {'Nodes':>8} {'GPUs':>8} {'CPUs':>8}")
            print("-" * 40)
            total_gpus = 0
            for user, stats in sorted_users:
                print(
                    f"{user:<15} {stats['nodes']:>8} {stats['gpus']:>8} {stats['cpus']:>8}"
                )
                total_gpus += stats["gpus"]
            print("-" * 40)
            print(
                f"{'TOTAL':<15} {sum(s['nodes'] for s in user_stats.values()):>8} "
                f"{total_gpus:>8} {sum(s['cpus'] for s in user_stats.values()):>8}"
            )
        else:
            print("No running jobs found.")

        # Show running jobs with commands
        running_jobs = self.get_running_jobs_with_commands()
        if running_jobs:
            print("\n📊 RUNNING JOBS DETAIL")
            print("-" * 140)
            print(
                f"{'JobID':<10} {'User':<10} {'Nodes':<6} {'GPUs':<6} {'Time':<12} {'Command':<30} {'WorkDir':<30} {'Output':<25} {'⏰ >3d':<5}"
            )
            print("-" * 140)
            total_nodes = 0
            for job in running_jobs:
                gpus = job["nodes"] * 8
                total_nodes += job["nodes"]

                # Check if job is long-running (>3 days)
                time_str = job["time"]
                long_running = ""
                if "-" in time_str:
                    days = int(time_str.split("-")[0])
                    if days >= 3:
                        long_running = "⏰"

                # Trim long strings with ... in the middle
                command = self.trim_middle(job["command"], 28)
                workdir = self.trim_middle(job["workdir"], 28)
                stdout = self.trim_middle(job["stdout"], 23)

                print(
                    f"{job['id']:<10} {job['user']:<10} {job['nodes']:<6} {gpus:<6} {job['time']:<12} {command:<30} {workdir:<30} {stdout:<25} {long_running:<5}"
                )
            print("-" * 140)
            total_gpus = total_nodes * 8
            print(
                f"{'TOTAL':<10} {len(running_jobs)} jobs{'':<6} {total_nodes:<6} {total_gpus:<6}"
            )

        # Problem Nodes
        print("\n⚠️  DRAINED NODES (Dev, or Down if no Reason)")
        print("-" * 80)
        drained = self.get_drained_nodes()
        if drained:
            print(f"Total: {len(drained)} problem nodes")
            print("-" * 80)
            print(f"{'Reason':<50} {'Timestamp':<20} {'State':<10} {'Nodes'}")
            print("-" * 80)
            for node_info in drained:  # Show all problem nodes
                print(node_info)
        else:
            print("No problem nodes detected")

        print("\n" + "=" * 80)
        print("END OF REPORT")
        print("=" * 80)

    def export_json(self, filename=None):
        """Export monitoring data to JSON."""
        if not filename:
            filename = f"cluster_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"

        data = {
            "timestamp": self.timestamp,
            "node_states": dict(self.get_node_summary()[0]),
            "user_stats": self.get_gpu_usage_by_user(),
            "idle_nodes_count": len(self.get_idle_nodes()),
            "queue_summary": self.get_queue_summary(),
            "problem_nodes_count": len(self.get_drained_nodes()),
        }

        with open(filename, "w") as f:
            json.dump(data, f, indent=2)

        return filename


def main():
    parser = argparse.ArgumentParser(description="SLURM Cluster Monitoring Tool")
    parser.add_argument("-v", "--verbose", action="store_true", help="Verbose output")
    parser.add_argument("-j", "--json", help="Export to JSON file")
    parser.add_argument(
        "-c",
        "--continuous",
        type=int,
        metavar="SECONDS",
        help="Run continuously with specified interval",
    )

    args = parser.parse_args()

    monitor = ClusterMonitor(verbose=args.verbose)

    if args.continuous:
        import time

        try:
            while True:
                # Clear screen
                print("\033[2J\033[H")
                monitor.print_report()

                if args.json:
                    filename = monitor.export_json(args.json)
                    print(f"\n📁 Report exported to: {filename}")

                print(
                    f"\n🔄 Refreshing in {args.continuous} seconds... (Ctrl+C to stop)"
                )
                time.sleep(args.continuous)
        except KeyboardInterrupt:
            print("\n\nMonitoring stopped.")
    else:
        monitor.print_report()

        if args.json:
            filename = monitor.export_json(args.json)
            print(f"\n📁 Report exported to: {filename}")


if __name__ == "__main__":
    main()