#!/usr/bin/env python3
"""
SLURM Job History Viewer - Shows completed and past jobs
"""

import subprocess
import argparse
from datetime import datetime, timedelta


def run_command(cmd):
    """Execute a shell command and return output."""
    try:
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        return result.stdout
    except Exception as e:
        print(f"Error running command: {e}")
        return ""


def get_job_history(user=None, days=1, state=None):
    """Get job history using sacct."""
    # Calculate start date
    start_date = (datetime.now() - timedelta(days=days)).strftime("%Y-%m-%d")

    # Build sacct command
    cmd = f"sacct -S {start_date} --format=JobID,JobName%30,User%15,State%20,Start,End,Elapsed,AllocCPUS,AllocNodes,ExitCode"

    if user:
        cmd += f" -u {user}"
    else:
        # Show all users when no specific user is requested
        cmd += " -a"

    if state:
        cmd += f" --state={state}"

    # Add -X flag to show only main job records (no steps/batch/extern)
    cmd += " -X"

    output = run_command(cmd)
    return output


def get_detailed_job_info(job_id):
    """Get detailed information about a specific job."""
    cmd = f"sacct -j {job_id} --format=JobID,JobName%30,User%15,State%20,Start,End,Elapsed,AllocCPUS,AllocNodes,ExitCode,NodeList%30,WorkDir%50"
    output = run_command(cmd)
    return output


def get_user_summary(days=7):
    """Get summary of jobs per user."""
    start_date = (datetime.now() - timedelta(days=days)).strftime("%Y-%m-%d")

    cmd = f"""sacct -S {start_date} -o User,State -n | awk '
    {{
        if ($2 == "COMPLETED") completed[$1]++
        else if ($2 == "FAILED") failed[$1]++
        else if ($2 == "CANCELLED") cancelled[$1]++
        else if ($2 == "TIMEOUT") timeout[$1]++
        total[$1]++
    }}
    END {{
        printf "%-15s %10s %10s %10s %10s %10s\\n", "User", "Total", "Completed", "Failed", "Cancelled", "Timeout"
        print "-------------------------------------------------------------------------------"
        for (user in total) {{
            if (user != "") {{
                printf "%-15s %10d %10d %10d %10d %10d\\n", 
                    user, total[user], 
                    completed[user]+0, 
                    failed[user]+0, 
                    cancelled[user]+0, 
                    timeout[user]+0
            }}
        }}
    }}'"""

    output = run_command(cmd)
    return output


def get_recent_failures(hours=24):
    """Get recently failed jobs."""
    # sacct command for failed jobs
    cmd = f"""sacct -S now-{hours}hours --state=FAILED,TIMEOUT,NODE_FAIL,OUT_OF_MEMORY \
        --format=JobID,User%15,JobName%30,State%20,Start,End,ExitCode,NodeList%20"""

    output = run_command(cmd)
    return output


def show_job_efficiency(job_id):
    """Show efficiency metrics for a completed job."""
    cmd = f"seff {job_id} 2>/dev/null"
    output = run_command(cmd)
    if not output:
        # Fallback to sacct if seff not available
        cmd = (
            f"sacct -j {job_id} --format=JobID,AllocCPUS,ReqMem,Elapsed,State,ExitCode"
        )
        output = run_command(cmd)
    return output


def main():
    parser = argparse.ArgumentParser(description="SLURM Job History Viewer")
    parser.add_argument(
        "-d",
        "--days",
        type=int,
        default=1,
        help="Number of days of history to show (default: 1)",
    )
    parser.add_argument("-u", "--user", help="Filter by specific user")
    parser.add_argument(
        "-s", "--state", help="Filter by job state (e.g., COMPLETED, FAILED, CANCELLED)"
    )
    parser.add_argument("-j", "--job", help="Show details for specific job ID")
    parser.add_argument(
        "--summary", action="store_true", help="Show user summary for past week"
    )
    parser.add_argument(
        "--failures",
        action="store_true",
        help="Show recent job failures (last 24 hours)",
    )
    parser.add_argument("--efficiency", help="Show efficiency for specific job ID")

    args = parser.parse_args()

    print("=" * 80)
    print(f"SLURM JOB HISTORY - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 80)

    if args.efficiency:
        print(f"\n📊 EFFICIENCY METRICS FOR JOB {args.efficiency}")
        print("-" * 40)
        print(show_job_efficiency(args.efficiency))

    elif args.job:
        print(f"\n📋 DETAILED INFO FOR JOB {args.job}")
        print("-" * 40)
        print(get_detailed_job_info(args.job))

    elif args.summary:
        print("\n👥 USER JOB SUMMARY (Past 7 Days)")
        print("-" * 40)
        print(get_user_summary())

    elif args.failures:
        print("\n⚠️  RECENT JOB FAILURES (Last 24 Hours)")
        print("-" * 40)
        failures = get_recent_failures()
        if failures:
            print(failures)
        else:
            print("No failures in the last 24 hours")

    else:
        # Default: show job history
        state_filter = f" ({args.state})" if args.state else ""
        user_filter = f" for user {args.user}" if args.user else ""
        print(f"\n📜 JOB HISTORY - Past {args.days} day(s){user_filter}{state_filter}")
        print("-" * 80)
        history = get_job_history(user=args.user, days=args.days, state=args.state)
        print(history)

        # Add summary stats
        if not args.state:
            print("\n📊 QUICK STATS")
            print("-" * 40)

            # Count jobs by state
            for state in ["COMPLETED", "FAILED", "CANCELLED", "RUNNING", "PENDING"]:
                cmd = f"sacct -S now-{args.days}days --state={state} -n -X"  # Added -X flag
                if args.user:
                    cmd += f" -u {args.user}"
                else:
                    cmd += " -a"  # Show all users
                cmd += " | wc -l"
                count = run_command(cmd).strip()
                if count and count != "0":
                    print(f"{state:15}: {count} jobs")


if __name__ == "__main__":
    main()