#!/usr/bin/env python3
"""
SLURM Jobs Detail Viewer - Shows comprehensive details for all running jobs
"""

import subprocess
import argparse
from datetime import datetime
import json


class JobDetailViewer:
    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 parse_scontrol_output(self, output):
        """Parse scontrol show job output into dictionary."""
        job_details = {}

        # Split into key=value pairs
        for line in output.split("\n"):
            # Each line can have multiple key=value pairs separated by spaces
            pairs = line.strip().split()
            for pair in pairs:
                if "=" in pair:
                    key, value = pair.split("=", 1)
                    job_details[key] = value

        return job_details

    def get_all_running_jobs(self):
        """Get list of all running job IDs."""
        cmd = "squeue -h -t RUNNING -o '%i'"
        output = self.run_command(cmd)

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

        return job_ids

    def get_job_details(self, job_id):
        """Get comprehensive details for a specific job."""
        cmd = f"scontrol show job {job_id}"
        output = self.run_command(cmd)

        if not output or "Invalid job id" in output:
            return None

        return self.parse_scontrol_output(output)

    def format_job_details(self, job_id, details):
        """Format job details for display."""
        if not details:
            return f"Job {job_id}: Unable to retrieve details"

        # Extract key fields with defaults
        formatted = []
        formatted.append(f"\n{'='*80}")
        formatted.append(f"JOB ID: {details.get('JobId', job_id)}")
        formatted.append(f"{'='*80}")

        # Basic Info
        formatted.append("\n📋 BASIC INFORMATION:")
        formatted.append(f"  Job Name:        {details.get('JobName', 'N/A')}")
        formatted.append(
            f"  User:            {details.get('UserId', 'N/A').split('(')[0] if 'UserId' in details else 'N/A'}"
        )
        formatted.append(f"  State:           {details.get('JobState', 'N/A')}")
        formatted.append(f"  Priority:        {details.get('Priority', 'N/A')}")
        formatted.append(f"  Partition:       {details.get('Partition', 'N/A')}")
        formatted.append(f"  QOS:             {details.get('QOS', 'N/A')}")

        # Time Info
        formatted.append("\n⏰ TIME INFORMATION:")
        formatted.append(f"  Submit Time:     {details.get('SubmitTime', 'N/A')}")
        formatted.append(f"  Start Time:      {details.get('StartTime', 'N/A')}")
        formatted.append(f"  Run Time:        {details.get('RunTime', 'N/A')}")
        formatted.append(f"  Time Limit:      {details.get('TimeLimit', 'N/A')}")
        formatted.append(f"  End Time:        {details.get('EndTime', 'N/A')}")

        # Resource Allocation
        formatted.append("\n💻 RESOURCE ALLOCATION:")
        formatted.append(f"  Nodes:           {details.get('NumNodes', 'N/A')}")
        formatted.append(f"  Node List:       {details.get('NodeList', 'N/A')}")
        formatted.append(f"  CPUs Total:      {details.get('NumCPUs', 'N/A')}")
        formatted.append(f"  CPUs Per Task:   {details.get('CPUs/Task', 'N/A')}")
        formatted.append(f"  Tasks:           {details.get('NumTasks', 'N/A')}")
        formatted.append(f"  Tasks Per Node:  {details.get('TasksPerNode', 'N/A')}")

        # Memory
        formatted.append(f"  Memory Per Node: {details.get('MinMemoryNode', 'N/A')}")
        formatted.append(f"  Memory Per CPU:  {details.get('MinMemoryCPU', 'N/A')}")

        # GPU Information
        if "GRES" in details or "Gres" in details:
            gres = details.get("GRES", details.get("Gres", "N/A"))
            formatted.append(f"  GPUs (GRES):     {gres}")

        # Paths and Files
        formatted.append("\n📁 PATHS AND FILES:")
        formatted.append(f"  Working Dir:     {details.get('WorkDir', 'N/A')}")
        formatted.append(f"  Command:         {details.get('Command', 'N/A')}")
        formatted.append(f"  StdOut:          {details.get('StdOut', 'N/A')}")
        formatted.append(f"  StdErr:          {details.get('StdErr', 'N/A')}")
        formatted.append(f"  Batch Script:    {details.get('BatchScript', 'N/A')}")

        # Environment
        formatted.append("\n🌍 ENVIRONMENT:")
        formatted.append(f"  Account:         {details.get('Account', 'N/A')}")
        formatted.append(f"  Reservation:     {details.get('Reservation', 'N/A')}")
        formatted.append(f"  Features:        {details.get('Features', 'N/A')}")
        formatted.append(f"  Licenses:        {details.get('Licenses', 'N/A')}")

        # Performance/Status
        formatted.append("\n📊 PERFORMANCE/STATUS:")
        formatted.append(f"  Exit Code:       {details.get('ExitCode', 'N/A')}")
        formatted.append(f"  Reason:          {details.get('Reason', 'N/A')}")
        formatted.append(f"  Dependency:      {details.get('Dependency', 'N/A')}")

        # Advanced Settings
        if self.verbose:
            formatted.append("\n⚙️  ADVANCED SETTINGS:")
            formatted.append(f"  Nice:            {details.get('Nice', 'N/A')}")
            formatted.append(f"  Requeue:         {details.get('Requeue', 'N/A')}")
            formatted.append(f"  Restarts:        {details.get('Restarts', 'N/A')}")
            formatted.append(f"  Array Job ID:    {details.get('ArrayJobId', 'N/A')}")
            formatted.append(f"  Array Task ID:   {details.get('ArrayTaskId', 'N/A')}")
            formatted.append(f"  Batch Host:      {details.get('BatchHost', 'N/A')}")
            formatted.append(f"  Sockets/Board:   {details.get('Sockets', 'N/A')}")
            formatted.append(
                f"  Cores/Socket:    {details.get('CoresPerSocket', 'N/A')}"
            )
            formatted.append(
                f"  Threads/Core:    {details.get('ThreadsPerCore', 'N/A')}"
            )

        return "\n".join(formatted)

    def print_summary_table(self, jobs_data):
        """Print a summary table of all jobs."""
        print("\n" + "=" * 120)
        print("JOBS SUMMARY TABLE")
        print("=" * 120)

        # Header
        print(
            f"{'JobID':<10} {'User':<10} {'Name':<20} {'State':<12} {'Nodes':<6} {'CPUs':<6} {'Time':<12} {'Partition':<12}"
        )
        print("-" * 120)

        for job_id, details in jobs_data.items():
            if details:
                user = (
                    details.get("UserId", "N/A").split("(")[0]
                    if "UserId" in details
                    else "N/A"
                )
                name = details.get("JobName", "N/A")[:18]  # Truncate long names
                state = details.get("JobState", "N/A")
                nodes = details.get("NumNodes", "N/A")
                cpus = details.get("NumCPUs", "N/A")
                runtime = details.get("RunTime", "N/A")
                partition = details.get("Partition", "N/A")

                print(
                    f"{job_id:<10} {user:<10} {name:<20} {state:<12} {nodes:<6} {cpus:<6} {runtime:<12} {partition:<12}"
                )

        print("-" * 120)
        print(f"Total Jobs: {len(jobs_data)}")

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

        # Convert to serializable format
        export_data = {
            "timestamp": self.timestamp,
            "total_jobs": len(jobs_data),
            "jobs": jobs_data,
        }

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

        return filename

    def print_report(self, job_filter=None, user_filter=None):
        """Print comprehensive job details report."""
        print("=" * 80)
        print(f"SLURM JOB DETAILS REPORT - {self.timestamp}")
        print("=" * 80)

        # Get job list
        job_ids = self.get_all_running_jobs()

        if not job_ids:
            print("\nNo running jobs found.")
            return {}

        print(f"\nFound {len(job_ids)} running job(s)")

        # Collect all job details
        jobs_data = {}
        for job_id in job_ids:
            if job_filter and job_id != job_filter:
                continue

            details = self.get_job_details(job_id)

            # Apply user filter if specified
            if user_filter and details:
                user = details.get("UserId", "").split("(")[0]
                if user != user_filter:
                    continue

            jobs_data[job_id] = details

        # Print summary table first
        if len(jobs_data) > 1:  # Only show table if multiple jobs
            self.print_summary_table(jobs_data)

        # Print detailed info for each job
        print("\n" + "=" * 80)
        print("DETAILED JOB INFORMATION")
        print("=" * 80)

        for job_id, details in sorted(jobs_data.items()):
            print(self.format_job_details(job_id, details))

        return jobs_data


def main():
    parser = argparse.ArgumentParser(description="SLURM Job Details Viewer")
    parser.add_argument("-j", "--job", help="Show details for specific job ID")
    parser.add_argument("-u", "--user", help="Filter by specific user")
    parser.add_argument(
        "-v", "--verbose", action="store_true", help="Show additional details"
    )
    parser.add_argument("--json", help="Export to JSON file")
    parser.add_argument(
        "-s", "--summary", action="store_true", help="Show summary table only"
    )

    args = parser.parse_args()

    viewer = JobDetailViewer(verbose=args.verbose)

    if args.summary:
        # Just show the summary table
        job_ids = viewer.get_all_running_jobs()
        jobs_data = {}
        for job_id in job_ids:
            details = viewer.get_job_details(job_id)
            if args.user and details:
                user = details.get("UserId", "").split("(")[0]
                if user != args.user:
                    continue
            jobs_data[job_id] = details

        viewer.print_summary_table(jobs_data)
    else:
        # Full report
        jobs_data = viewer.print_report(job_filter=args.job, user_filter=args.user)

        if args.json and jobs_data:
            filename = viewer.export_json(jobs_data, args.json)
            print(f"\n📁 Details exported to: {filename}")


if __name__ == "__main__":
    main()