#!/usr/bin/python3
"""
Slurm Quota Management Tool

Copyright (c) 2025 Rackslab
SPDX-License-Identifier: MIT

This tool tracks user resource consumption by maintaining a SQLite database.

It has the following subcommands:
- charge: Update user/account resource consumption. This is designed to be executed by
  Slurm as a JobCompType=jobcomp/script.
- stats: Show resource statistics for users and accounts.
- adjust: Adjust consumed CPU/GPU minutes on a user/account (restricted to root only).
- user-quota: Set quota for a user (restricted to root only).
- account-quota: Set quota for an account (restricted to root only).
- user-gpu-quota: Set GPU quota for a user (restricted to root only).
- account-gpu-quota: Set GPU quota for an account (restricted to root only).
- default-quotas: Show default quotas for newly auto-created users/accounts.
- set-default-quotas: Set default quotas for newly auto-created users/accounts
  (restricted to root only).
- gpu-factors: Show configured GPU charging factors.
- set-gpu-factor: Set charging factor for a GPU type (restricted to root only).
- prune: Prune selected preallocs/users/accounts records (root only).
- serve: Serve HTTP JSON API for stats (socket-activated by systemd; falls back to host/port).
"""

import argparse
import sqlite3
import sys
import os
import logging
import pwd
import re
import subprocess
import json
import socket
import time
from datetime import datetime, timezone
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, HTTPServer
from urllib.request import urlopen, Request
from urllib.error import URLError
from urllib.parse import urlencode, urljoin, urlparse, parse_qs
from typing import Tuple, Union, Optional, Any, Dict, List

# Database configuration
DB_PATH = "/var/lib/state/slurm-quota/slurm-quota.db"
APP_VERSION = "2.0.0"
DEFAULT_QUOTA_SETTINGS = {
    "default_user_quota_cpu_minutes": -1,
    "default_user_quota_gpu_minutes": -1,
    "default_account_quota_cpu_minutes": -1,
    "default_account_quota_gpu_minutes": -1,
}


# Configure logging
def setup_logging(debug: bool = False, quiet: bool = False) -> None:
    """
    Setup logging configuration.

    Args:
        debug: If True, set logging level to DEBUG
        quiet: If True, set logging level to WARNING
    """
    if debug:
        level = logging.DEBUG
    elif quiet:
        level = logging.WARNING
    else:
        level = logging.INFO
    logging.basicConfig(
        level=level,
        format="%(asctime)s - %(levelname)s - %(message)s",
        handlers=[
            logging.StreamHandler(sys.stderr),
        ],
        force=True,
    )


logger = logging.getLogger(__name__)


def get_default_quota_settings() -> Dict[str, int]:
    """
    Return configured default quotas used for new users/accounts.

    Missing or invalid values fall back to -1 (unlimited) for robustness.
    """
    defaults = dict(DEFAULT_QUOTA_SETTINGS)

    if not os.path.exists(DB_PATH):
        return defaults

    try:
        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()
            cursor.execute(
                "SELECT key, value FROM settings WHERE key IN (?, ?, ?, ?)",
                tuple(DEFAULT_QUOTA_SETTINGS.keys()),
            )
            for key, value in cursor.fetchall():
                if key in defaults:
                    try:
                        defaults[key] = int(value)
                    except (TypeError, ValueError):
                        logger.warning(
                            "Invalid default quota setting for %s=%r; using %d",
                            key,
                            value,
                            DEFAULT_QUOTA_SETTINGS[key],
                        )
    except sqlite3.Error as e:
        logger.warning("Failed to load default quota settings from database: %s", e)

    return defaults


def set_default_quota_settings(
    default_user_cpu: Optional[int],
    default_user_gpu: Optional[int],
    default_account_cpu: Optional[int],
    default_account_gpu: Optional[int],
) -> None:
    """
    Set one or more default quotas used for newly auto-created entities.
    """
    updates: Dict[str, int] = {}
    if default_user_cpu is not None:
        updates["default_user_quota_cpu_minutes"] = default_user_cpu
    if default_user_gpu is not None:
        updates["default_user_quota_gpu_minutes"] = default_user_gpu
    if default_account_cpu is not None:
        updates["default_account_quota_cpu_minutes"] = default_account_cpu
    if default_account_gpu is not None:
        updates["default_account_quota_gpu_minutes"] = default_account_gpu

    if not updates:
        raise ValueError("No default quotas provided")

    with sqlite3.connect(DB_PATH) as conn:
        cursor = conn.cursor()
        cursor.executemany(
            """
            INSERT INTO settings (key, value)
            VALUES (?, ?)
            ON CONFLICT(key) DO UPDATE SET value = excluded.value
            """,
            [(key, str(value)) for key, value in updates.items()],
        )
        conn.commit()


def load_gpu_factors() -> Dict[str, float]:
    """
    Load GPU type charging factors from the SQLite database.

    The factors are stored in the ``gpu_factors`` table with:
      - gpu_type TEXT PRIMARY KEY
      - factor   REAL NOT NULL

    A special gpu_type value of ``default`` controls the default factor used
    when no explicit entry exists for a given GPU type.

    Returns:
        Dictionary mapping GPU type to factor. The key ``"__default__"`` holds
        the default factor (1.0 when not configured).
    """
    factors: Dict[str, float] = {}
    default_factor = 1.0

    if not os.path.exists(DB_PATH):
        # No database yet – use built-in defaults
        factors["__default__"] = default_factor
        return factors

    try:
        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT gpu_type, factor FROM gpu_factors")
            rows = cursor.fetchall()

            for gpu_type, factor in rows:
                try:
                    factor_f = float(factor)
                except (TypeError, ValueError):
                    logger.warning(
                        "Invalid GPU factor value for type %r: %r; skipping",
                        gpu_type,
                        factor,
                    )
                    continue

                if gpu_type == "default":
                    if factor_f < 0:
                        logger.warning(
                            "Invalid default GPU factor %r; keeping 1.0", factor_f
                        )
                    else:
                        default_factor = factor_f
                else:
                    if factor_f < 0:
                        logger.warning(
                            "Invalid GPU factor %r for type %r; skipping",
                            factor_f,
                            gpu_type,
                        )
                        continue
                    factors[str(gpu_type)] = factor_f
    except sqlite3.Error as e:
        logger.warning("Failed to load GPU factors from database: %s", e)

    factors["__default__"] = default_factor
    return factors


def parse_alloc_tres(alloc_tres: str) -> Dict[str, int]:
    """
    Parse AllocTRES field from sacct to extract GPU allocations.

    Args:
        alloc_tres: AllocTRES string, e.g., "billing=1,cpu=1,gres/gpu:h100=2,gres/gpu:h200=1,mem=512M,node=1"

    Returns:
        Dictionary mapping GPU type to count (e.g., {"h100": 2, "h200": 1})
    """
    gpu_counts: Dict[str, int] = {}

    if not alloc_tres:
        return gpu_counts

    # Split by comma and look for gres/gpu:TYPE=N patterns
    for item in alloc_tres.split(","):
        item = item.strip()
        # Match pattern: gres/gpu:TYPE=N or gres/gpu:TYPE=N
        match = re.match(r"gres/gpu:([^=]+)=(\d+)", item)
        if match:
            gpu_type = match.group(1).strip()
            count = int(match.group(2))
            gpu_counts[gpu_type] = gpu_counts.get(gpu_type, 0) + count

    return gpu_counts


def calculate_consumed_gpu_minutes(
    gpu_counts: Dict[str, int], job_duration_minutes: int, factors: Dict[str, float]
) -> int:
    """
    Calculate consumed GPU minutes based on GPU allocations, duration, and factors.

    Args:
        gpu_counts: Dictionary mapping GPU type to count
        job_duration_minutes: Job duration in minutes
        factors: Dictionary mapping GPU type to factor (from load_gpu_factors)

    Returns:
        Total GPU minutes consumed (integer)
    """
    default_factor = factors.get("__default__", 1.0)
    total_gpu_minutes = 0.0

    for gpu_type, count in gpu_counts.items():
        factor = factors.get(gpu_type, default_factor)
        total_gpu_minutes += count * job_duration_minutes * factor

    return int(total_gpu_minutes)


def format_timestamp_with_timezone(timestamp_str: Optional[str]) -> str:
    """
    Format a timestamp string to display with local timezone.

    Args:
        timestamp_str: UTC timestamp string from REST API in ISO format

    Returns:
        Formatted timestamp string with local timezone, or "N/A" if invalid
    """
    if not timestamp_str:
        return "N/A"

    try:
        # ISO format from API without timezone - assume UTC
        dt = datetime.fromisoformat(timestamp_str)
        # Replace timezone info to treat as UTC
        dt = dt.replace(tzinfo=timezone.utc)

        # Convert to local timezone
        local_dt = dt.astimezone()
        # Format as YYYY-MM-DD HH:MM:SS TZ
        return local_dt.strftime("%Y-%m-%d %H:%M:%S %Z")
    except (ValueError, TypeError) as e:
        logger.debug(f"Failed to parse timestamp '{timestamp_str}': {e}")
        return "N/A"


def parse_slurm_time_format(time_str: str) -> Union[int, None]:
    """
    Parse Slurm time format string and convert to minutes.

    Slurm formats time as:
    - "UNLIMITED" for infinite time
    - "INVALID" for invalid time
    - "DD-HH:MM:SS" for durations with days
    - "HH:MM:SS" for durations without days

    Args:
        time_str: Time string from Slurm LIMIT environment variable

    Returns:
        Number of minutes, or None for UNLIMITED/INVALID cases

    Raises:
        ValueError: If the time format is invalid
    """
    if not time_str:
        return 0

    time_str = time_str.strip()

    # Handle empty string after stripping
    if not time_str:
        return 0

    # Handle special cases
    if time_str == "UNLIMITED":
        return None  # None represents unlimited
    elif time_str == "INVALID":
        raise ValueError("Invalid time format")

    # Pattern for DD-HH:MM:SS format
    days_pattern = r"^(\d+)-(\d{2}):(\d{2}):(\d{2})$"
    # Pattern for HH:MM:SS format
    hours_pattern = r"^(\d{2}):(\d{2}):(\d{2})$"

    days_match = re.match(days_pattern, time_str)
    if days_match:
        days, hours, minutes, seconds = map(int, days_match.groups())
        if days < 0 or hours < 0 or minutes < 0 or seconds < 0:
            raise ValueError("Negative time values not allowed")
        if hours >= 24 or minutes >= 60 or seconds >= 60:
            raise ValueError(
                "Invalid time values: hours must be < 24, minutes and seconds must be < 60"
            )
        return days * 1440 + hours * 60 + minutes

    hours_match = re.match(hours_pattern, time_str)
    if hours_match:
        hours, minutes, seconds = map(int, hours_match.groups())
        if hours < 0 or minutes < 0 or seconds < 0:
            raise ValueError("Negative time values not allowed")
        if hours >= 24 or minutes >= 60 or seconds >= 60:
            raise ValueError(
                "Invalid time values: hours must be < 24, minutes and seconds must be < 60"
            )
        return hours * 60 + minutes

    # If no pattern matches, try to parse as integer (fallback for backward compatibility)
    try:
        return int(time_str)
    except ValueError:
        raise ValueError(f"Unrecognized time format: {time_str}")


def parse_signed_int(value: str) -> int:
    """
    Parse an explicitly signed integer for CLI arguments.

    Args:
        value: Raw argument value

    Returns:
        Parsed integer value

    Raises:
        argparse.ArgumentTypeError: If value does not include an explicit sign
    """
    if not re.fullmatch(r"[+-]\d+", value):
        raise argparse.ArgumentTypeError(
            "must be an explicitly signed integer (for example: +30 or -30)"
        )
    return int(value)


def get_current_user() -> str:
    """
    Get the current user name based on the process UID.

    Returns:
        The username of the current process owner
    """
    current_uid = os.getuid()
    try:
        return pwd.getpwuid(current_uid).pw_name
    except KeyError:
        logger.error(f"Unable to get user name for UID {current_uid}")
        raise


# Unfortunately, Slurm does not provide a way to get the admin_comment field in
# jobcomp script environment variables for a job, so we need to use sacct to get it.
def get_job_info_from_sacct(job_id: str) -> Tuple[Optional[str], Optional[str]]:
    """
    Get the admin_comment and AllocTRES fields for a Slurm job.

    Args:
        job_id: The Slurm job ID

    Returns:
        Tuple of (admin_comment, alloc_tres) or (None, None) if not found/error
    """
    try:
        result = subprocess.run(
            [
                "sacct",
                "--noheader",
                "--allocations",
                "--job",
                job_id,
                "--parsable2",
                "--format",
                "admincomment,alloctres",
            ],
            capture_output=True,
            text=True,
            check=True,
        )
        output = result.stdout.strip()
        if not output:
            return None, None

        # Parsable2 format uses | as delimiter
        parts = output.split("|", 1)
        admin_comment = parts[0].strip() if len(parts) > 0 else ""
        alloc_tres = parts[1].strip() if len(parts) > 1 else ""

        # Handle (null) case
        if admin_comment == "(null)" or not admin_comment:
            admin_comment = None
        if alloc_tres == "(null)" or not alloc_tres:
            alloc_tres = None

        return admin_comment, alloc_tres
    except (subprocess.CalledProcessError, FileNotFoundError) as e:
        logger.warning(f"Failed to get job info from sacct for job {job_id}: {e}")
        return None, None


def get_user_accounts(username: str) -> set[str]:
    """
    Return the set of Slurm accounts the given user belongs to using sacctmgr.

    Args:
        username: The Slurm username

    Returns:
        A set of account names (may be empty when none or on error)
    """
    try:
        result = subprocess.run(
            [
                "sacctmgr",
                "list",
                "associations",
                "where",
                f"user={username}",
                "format=account",
                "-P",
                "-n",
            ],
            capture_output=True,
            text=True,
            check=True,
        )
        accounts = {line.strip() for line in result.stdout.splitlines() if line.strip()}
        logger.debug(f"User {username} belongs to accounts: {accounts}")
        return accounts
    except (subprocess.CalledProcessError, FileNotFoundError) as e:
        logger.warning(f"Failed to get accounts for user {username} via sacctmgr: {e}")
        return set()


def init_database() -> None:
    """
    Initialize the SQLite database only if it doesn't exist yet. Set permissions on
    first creation.
    """
    try:
        # Ensure the directory exists
        os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)

        # If the database file already exists, do nothing
        if os.path.exists(DB_PATH):
            return

        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()

            # Create users table to track resource usage
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS users (
                    username TEXT PRIMARY KEY,
                    total_consumed_cpu_minutes INTEGER DEFAULT 0,
                    quota_cpu_minutes INTEGER DEFAULT -1,
                    total_consumed_gpu_minutes INTEGER DEFAULT 0,
                    quota_gpu_minutes INTEGER DEFAULT -1,
                    last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)

            # Create accounts table to track resource usage
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS accounts (
                    account TEXT PRIMARY KEY,
                    total_consumed_cpu_minutes INTEGER DEFAULT 0,
                    quota_cpu_minutes INTEGER DEFAULT -1,
                    total_consumed_gpu_minutes INTEGER DEFAULT 0,
                    quota_gpu_minutes INTEGER DEFAULT -1,
                    last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)

            # Create jobs_preallocations table to track preallocated resources per job
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS jobs_preallocations (
                    job_uuid TEXT PRIMARY KEY,
                    username TEXT NOT NULL,
                    account TEXT NOT NULL,
                    preallocated_cpu_minutes INTEGER NOT NULL,
                    preallocated_gpu_minutes INTEGER DEFAULT 0,
                    array_size INTEGER DEFAULT 1,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (username) REFERENCES users (username),
                    FOREIGN KEY (account) REFERENCES accounts (account)
                )
            """)

            # Create gpu_factors table to store GPU charging factors
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS gpu_factors (
                    gpu_type TEXT PRIMARY KEY,
                    factor REAL NOT NULL
                )
            """)

            # Create settings table for runtime configuration values
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS settings (
                    key TEXT PRIMARY KEY,
                    value TEXT NOT NULL
                )
            """)
            cursor.executemany(
                """
                INSERT INTO settings (key, value)
                VALUES (?, ?)
                ON CONFLICT(key) DO NOTHING
                """,
                [(key, str(value)) for key, value in DEFAULT_QUOTA_SETTINGS.items()],
            )

            conn.commit()
            logger.info("Database initialized successfully")

        # Set permissions on the database file
        set_database_permissions()

    except sqlite3.Error as e:
        logger.error(f"Database initialization failed: {e}")
        raise
    except OSError as e:
        logger.error(f"Failed to create database directory: {e}")
        raise


def set_database_permissions() -> None:
    """
    Set database file permissions to allow root and slurm users to modify it and
    others to read it.
    """
    try:
        # If root, give database file to slurm user
        if get_current_user() == "root":
            logger.info(f"Giving database file to slurm user for {DB_PATH}")
            slurm_pw = pwd.getpwnam("slurm")
            os.chown(DB_PATH, slurm_pw.pw_uid, slurm_pw.pw_gid)
        # Set permissions: owner (slurm) can read/write, others can read
        os.chmod(DB_PATH, 0o644)
        logger.info(f"Database permissions set for {DB_PATH}")
    except OSError as e:
        logger.error(f"Failed to set database permissions: {e}")
        raise


def collect_active_job_uuids() -> set[str]:
    """
    Collect active Slurm job UUIDs from squeue admin_comment field.

    Returns:
        A set of admin_comment values (UUIDs) for all jobs visible in the queue.
    """
    uuids: set[str] = set()
    try:
        result = subprocess.run(
            ["squeue", "--noheader", "--Format", "jobid,admin_comment:50"],
            capture_output=True,
            text=True,
            check=True,
        )
        for line in result.stdout.splitlines():
            line = line.strip()
            # squeeze whitespaces to get a single space between jobid and comment
            line = re.sub(r"\s+", " ", line)
            if not line:
                continue
            # Split on first whitespace to separate jobid and comment
            parts = line.split(None, 1)
            if len(parts) == 2:
                _jobid, comment = parts
                comment = comment.strip()
                if comment:
                    uuids.add(comment)
    except (subprocess.CalledProcessError, FileNotFoundError) as e:
        logger.warning(f"Failed to collect active job UUIDs via squeue: {e}")
    return uuids


def prune_resources(
    targets: set[str],
    dry_run: bool = False,
    user_filter: Optional[str] = None,
    account_filter: Optional[str] = None,
) -> Dict[str, int]:
    """
    Prune selected resource records from the database.

    Args:
        targets: Set of prune targets among {"preallocs", "users", "accounts"}.
        dry_run: If True, report counts without deleting rows.
        user_filter: Optional username to limit user pruning candidates.
        account_filter: Optional account to limit account pruning candidates.

    Returns:
        Mapping with deleted (or deletable in dry-run) rows per target.
    """
    counts = {"preallocs": 0, "users": 0, "accounts": 0}
    if not os.path.exists(DB_PATH):
        logger.info("Database not found; nothing to prune")
        return counts

    with sqlite3.connect(DB_PATH) as conn:
        conn.execute("PRAGMA foreign_keys = ON")
        cursor = conn.cursor()

        prealloc_rows: List[Tuple[Optional[str], Optional[str], Optional[str]]] = []
        orphan_prealloc_uuids: List[str] = []
        if "preallocs" in targets:
            cursor.execute(
                "SELECT job_uuid, username, account FROM jobs_preallocations"
            )
            prealloc_rows = cursor.fetchall()
            active_job_uuids = collect_active_job_uuids()
            orphan_prealloc_uuids = [
                str(job_uuid)
                for job_uuid, _username, _account in prealloc_rows
                if job_uuid and job_uuid not in active_job_uuids
            ]
            counts["preallocs"] = len(orphan_prealloc_uuids)

        users_to_delete: List[str] = []
        if "users" in targets:
            users_query = """
                SELECT username
                FROM users
                WHERE total_consumed_cpu_minutes = 0
                  AND total_consumed_gpu_minutes = 0
            """
            users_params: List[str] = []
            if user_filter:
                users_query += " AND username = ?"
                users_params.append(user_filter)
            cursor.execute(users_query, users_params)
            users_to_delete = [
                str(username) for (username,) in cursor.fetchall() if username
            ]
            for username in users_to_delete:
                logger.info("Eligible user for pruning: %s", username)
            counts["users"] = len(users_to_delete)

        accounts_to_delete: List[str] = []
        if "accounts" in targets:
            accounts_query = """
                SELECT account
                FROM accounts
                WHERE total_consumed_cpu_minutes = 0
                  AND total_consumed_gpu_minutes = 0
            """
            accounts_params: List[str] = []
            if account_filter:
                accounts_query += " AND account = ?"
                accounts_params.append(account_filter)
            cursor.execute(accounts_query, accounts_params)
            accounts_to_delete = [
                str(account) for (account,) in cursor.fetchall() if account
            ]
            for account in accounts_to_delete:
                logger.info("Eligible account for pruning: %s", account)
            counts["accounts"] = len(accounts_to_delete)

        if dry_run:
            return counts

        # Delete preallocations first so users/accounts can become eligible in the same run.
        if orphan_prealloc_uuids:
            cursor.executemany(
                "DELETE FROM jobs_preallocations WHERE job_uuid = ?",
                [(job_uuid,) for job_uuid in orphan_prealloc_uuids],
            )

        if users_to_delete:
            try:
                cursor.executemany(
                    "DELETE FROM users WHERE username = ?",
                    [(username,) for username in users_to_delete],
                )
            except sqlite3.IntegrityError as e:
                raise sqlite3.IntegrityError(
                    "Failed to prune users: one or more users are still referenced "
                    "by jobs preallocations"
                ) from e

        if accounts_to_delete:
            try:
                cursor.executemany(
                    "DELETE FROM accounts WHERE account = ?",
                    [(account,) for account in accounts_to_delete],
                )
            except sqlite3.IntegrityError as e:
                raise sqlite3.IntegrityError(
                    "Failed to prune accounts: one or more accounts are still "
                    "referenced by jobs preallocations"
                ) from e

        conn.commit()
        return counts


def get_job_info_from_environment() -> Tuple[str, str, int, str, Optional[str], int]:
    """
    Extract job information from Slurm environment variables.

    Returns:
        Tuple of (username, job_id, consumed_cpu_minutes, account, job_uuid, consumed_gpu_minutes)

    Raises:
        ValueError: If the job ID or username is not found
        ValueError: If the start or end time is not found
    """
    try:
        # Get job ID
        job_id = os.environ.get("JOBID", "unknown")

        # Check job is member of an array. If yes, compute the job ID as
        # ARRAYJOBID_ARRAYTASKID.
        array_job_id = int(os.environ.get("ARRAYJOBID", "0"))
        array_task_id = int(os.environ.get("ARRAYTASKID", "0"))
        if array_job_id != 0:
            job_id = f"{array_job_id}_{array_task_id}"

        # Get username from Slurm environment
        username = os.environ.get("USERNAME")
        if not username:
            raise ValueError(
                "USERNAME environment variable not found for job %s", job_id
            )

        # Get account from Slurm environment
        account = os.environ.get("ACCOUNT")
        if not account:
            raise ValueError(
                "ACCOUNT environment variable not found for job %s", job_id
            )

        # Calculate consumed CPU minutes
        # PROCS * (END - START) / 60
        job_procs = int(os.environ.get("PROCS", 1))
        job_start_time = int(os.environ.get("START", 0))
        job_end_time = int(os.environ.get("END", 0))
        # If START or END is not found, consider the job ran for 0 minutes
        if not job_start_time or not job_end_time:
            job_duration_minutes = 0
        else:
            job_duration_minutes = (job_end_time - job_start_time) // 60

        consumed_cpu_minutes = job_procs * job_duration_minutes

        # Get job UUID and AllocTRES from sacct
        job_uuid, alloc_tres = get_job_info_from_sacct(job_id)
        if not job_uuid:
            logger.warning(f"No job UUID found in admin_comment for job {job_id}")

        # Calculate consumed GPU minutes
        consumed_gpu_minutes = 0
        if alloc_tres:
            gpu_factors = load_gpu_factors()
            gpu_counts = parse_alloc_tres(alloc_tres)
            if gpu_counts:
                consumed_gpu_minutes = calculate_consumed_gpu_minutes(
                    gpu_counts, job_duration_minutes, gpu_factors
                )

        return (
            username,
            job_id,
            consumed_cpu_minutes,
            account,
            job_uuid,
            consumed_gpu_minutes,
        )

    except (ValueError, KeyError) as e:
        logger.error(f"Failed to extract job information: {e}")
        raise


def update_user_and_account_resources(
    username: str,
    account: str,
    consumed_cpu_minutes: int,
    job_uuid: Optional[str],
    consumed_gpu_minutes: int = 0,
) -> str:
    """
    Update user and account resource consumption in the database and adjust any
    preallocation associated with the job UUID.

    Args:
        username: The username of the job owner
        account: The account of the job
        consumed_cpu_minutes: Actual CPU minutes consumed by the job
        job_uuid: The job UUID from admin_comment
        consumed_gpu_minutes: Actual GPU minutes consumed by the job
    Returns:
        One of:
        - "removed": a preallocation row was fully removed
        - "decremented": array_size was decremented but row kept
        - "none": no preallocation was changed (including when job_uuid is None)
    """
    try:
        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()
            dq = get_default_quota_settings()
            default_user_cpu = dq["default_user_quota_cpu_minutes"]
            default_user_gpu = dq["default_user_quota_gpu_minutes"]
            default_account_cpu = dq["default_account_quota_cpu_minutes"]
            default_account_gpu = dq["default_account_quota_gpu_minutes"]

            # Update consumed CPU and GPU minutes for user
            cursor.execute(
                """
                INSERT INTO users (username, total_consumed_cpu_minutes, quota_cpu_minutes,
                    total_consumed_gpu_minutes, quota_gpu_minutes, last_updated)
                VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
                ON CONFLICT (username) DO
                UPDATE SET total_consumed_cpu_minutes = total_consumed_cpu_minutes + ?,
                           total_consumed_gpu_minutes = total_consumed_gpu_minutes + ?,
                           last_updated = CURRENT_TIMESTAMP
                """,
                (
                    username,
                    consumed_cpu_minutes,
                    default_user_cpu,
                    consumed_gpu_minutes,
                    default_user_gpu,
                    consumed_cpu_minutes,
                    consumed_gpu_minutes,
                ),
            )

            # Update consumed CPU and GPU minutes for account
            cursor.execute(
                """
                INSERT INTO accounts (account, total_consumed_cpu_minutes, quota_cpu_minutes,
                    total_consumed_gpu_minutes, quota_gpu_minutes, last_updated)
                VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
                ON CONFLICT (account) DO
                UPDATE SET total_consumed_cpu_minutes = total_consumed_cpu_minutes + ?,
                           total_consumed_gpu_minutes = total_consumed_gpu_minutes + ?,
                            last_updated = CURRENT_TIMESTAMP
                """,
                (
                    account,
                    consumed_cpu_minutes,
                    default_account_cpu,
                    consumed_gpu_minutes,
                    default_account_gpu,
                    consumed_cpu_minutes,
                    consumed_gpu_minutes,
                ),
            )

            # Update or remove preallocation if job_uuid exists
            if job_uuid:
                prealloc_status = "none"
                # Fetch current array_size for this job_uuid
                cursor.execute(
                    "SELECT array_size FROM jobs_preallocations WHERE job_uuid = ?",
                    (job_uuid,),
                )
                row = cursor.fetchone()
                if row is None:
                    logger.warning(
                        f"No preallocated resources found for job UUID: {job_uuid}"
                    )
                else:
                    current_array_size = int(row[0] or 0)
                    if current_array_size <= 1:
                        # Last element of the array: remove preallocation entirely
                        cursor.execute(
                            "DELETE FROM jobs_preallocations WHERE job_uuid = ?",
                            (job_uuid,),
                        )
                        if (cursor.rowcount or 0) > 0:
                            logger.debug(
                                "Removed preallocated resources for job UUID: %s",
                                job_uuid,
                            )
                            prealloc_status = "removed"
                    else:
                        # Decrement remaining array_size
                        new_array_size = current_array_size - 1
                        cursor.execute(
                            """
                            UPDATE jobs_preallocations
                            SET array_size = ?
                            WHERE job_uuid = ?
                            """,
                            (new_array_size, job_uuid),
                        )
                        logger.debug(
                            "Decremented preallocation array_size for job UUID %s: %d -> %d",
                            job_uuid,
                            current_array_size,
                            new_array_size,
                        )
                        prealloc_status = "decremented"
            else:
                prealloc_status = "none"

            conn.commit()

            return prealloc_status

    except sqlite3.Error as e:
        logger.error(f"Failed to update user resources: {e}")
        raise


def charge_command() -> None:
    """
    Execute the charge command to update user resource consumption.
    This command can only be executed by the slurm system user.
    """
    try:
        # Check if running as slurm user
        current_user = get_current_user()
        if current_user != "slurm":
            logger.error(
                f"Charge command can only be executed by slurm user, not by "
                f"{current_user}"
            )
            sys.exit(1)

        # Initialize database only if missing
        init_database()

        # Get job information from environment
        (
            username,
            job_id,
            consumed_cpu_minutes,
            account,
            job_uuid,
            consumed_gpu_minutes,
        ) = get_job_info_from_environment()

        # Update user and account resources (and any related preallocation)
        prealloc_status = update_user_and_account_resources(
            username, account, consumed_cpu_minutes, job_uuid, consumed_gpu_minutes
        )

        # Single summary info log on success
        logger.info(
            (
                f"charge: user={username} account={account} job_id={job_id} "
                f"+{consumed_cpu_minutes} CPUmins +{consumed_gpu_minutes} GPUmins "
                f"uuid={job_uuid or 'none'} prealloc_status={prealloc_status}"
            )
        )

    except Exception as e:
        logger.error(f"Charge command failed: {e}")
        sys.exit(1)


def set_user_quota(username: str, quota_cpu_minutes: int) -> None:
    """
    Set quota for a user in the database.

    Args:
        username: The username to set quota for
        quota_cpu_minutes: The quota in CPU minutes (-1 for unlimited)
    """
    try:
        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()
            dq = get_default_quota_settings()
            default_gpu = dq["default_user_quota_gpu_minutes"]

            # Insert or update user with quota
            cursor.execute(
                """
                INSERT INTO users (username, quota_cpu_minutes, quota_gpu_minutes, last_updated)
                VALUES (?, ?, ?, CURRENT_TIMESTAMP)
                ON CONFLICT (username) DO
                UPDATE SET quota_cpu_minutes = ?, last_updated = CURRENT_TIMESTAMP
                """,
                (username, quota_cpu_minutes, default_gpu, quota_cpu_minutes),
            )

            conn.commit()

    except sqlite3.Error as e:
        logger.error(f"Failed to set user quota: {e}")
        raise


def set_account_quota(account: str, quota_cpu_minutes: int) -> None:
    """
    Set quota for an account in the database.

    Args:
        account: The account to set quota for
        quota_cpu_minutes: The quota in CPU minutes (-1 for unlimited)
    """
    try:
        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()
            dq = get_default_quota_settings()
            default_gpu = dq["default_account_quota_gpu_minutes"]

            cursor.execute(
                """
                INSERT INTO accounts (account, quota_cpu_minutes, quota_gpu_minutes, last_updated)
                VALUES (?, ?, ?, CURRENT_TIMESTAMP)
                ON CONFLICT (account) DO
                UPDATE SET quota_cpu_minutes = ?, last_updated = CURRENT_TIMESTAMP
                """,
                (account, quota_cpu_minutes, default_gpu, quota_cpu_minutes),
            )

            conn.commit()

    except sqlite3.Error as e:
        logger.error(f"Failed to set account quota: {e}")
        raise


def set_gpu_factor(gpu_type: str, factor: float) -> None:
    """
    Set charging factor for a GPU type in the database.

    Args:
        gpu_type: GPU type name (use "default" for the default factor)
        factor: Charging factor (float, must be non-negative)
    """
    try:
        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()

            cursor.execute(
                """
                INSERT INTO gpu_factors (gpu_type, factor)
                VALUES (?, ?)
                ON CONFLICT(gpu_type) DO UPDATE SET factor = excluded.factor
                """,
                (gpu_type, float(factor)),
            )

            conn.commit()

    except sqlite3.Error as e:
        logger.error(f"Failed to set GPU factor: {e}")
        raise


def set_user_gpu_quota(username: str, quota_gpu_minutes: int) -> None:
    """
    Set GPU quota for a user in the database.

    Args:
        username: The username to set GPU quota for
        quota_gpu_minutes: The quota in GPU minutes (-1 for unlimited)
    """
    try:
        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()
            dq = get_default_quota_settings()
            default_cpu = dq["default_user_quota_cpu_minutes"]

            # Insert or update user with GPU quota
            cursor.execute(
                """
                INSERT INTO users (username, quota_cpu_minutes, quota_gpu_minutes, last_updated)
                VALUES (?, ?, ?, CURRENT_TIMESTAMP)
                ON CONFLICT (username) DO
                UPDATE SET quota_gpu_minutes = ?, last_updated = CURRENT_TIMESTAMP
                """,
                (username, default_cpu, quota_gpu_minutes, quota_gpu_minutes),
            )

            conn.commit()

    except sqlite3.Error as e:
        logger.error(f"Failed to set user GPU quota: {e}")
        raise


def set_account_gpu_quota(account: str, quota_gpu_minutes: int) -> None:
    """
    Set GPU quota for an account in the database.

    Args:
        account: The account to set GPU quota for
        quota_gpu_minutes: The quota in GPU minutes (-1 for unlimited)
    """
    try:
        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()
            dq = get_default_quota_settings()
            default_cpu = dq["default_account_quota_cpu_minutes"]

            cursor.execute(
                """
                INSERT INTO accounts (account, quota_cpu_minutes, quota_gpu_minutes, last_updated)
                VALUES (?, ?, ?, CURRENT_TIMESTAMP)
                ON CONFLICT (account) DO
                UPDATE SET quota_gpu_minutes = ?, last_updated = CURRENT_TIMESTAMP
                """,
                (account, default_cpu, quota_gpu_minutes, quota_gpu_minutes),
            )

            conn.commit()

    except sqlite3.Error as e:
        logger.error(f"Failed to set account GPU quota: {e}")
        raise


def adjust_consumed_minutes(
    target_type: str, target_name: str, resource: str, delta_minutes: int
) -> int:
    """
    Adjust consumed CPU/GPU minutes for a user or account.

    Args:
        target_type: "user" or "account"
        target_name: Username or account name
        resource: "cpu" or "gpu"
        delta_minutes: Signed minutes to add/subtract

    Returns:
        New consumed minutes value after update

    Raises:
        ValueError: If target type/resource is invalid or target row is missing
    """
    target_mapping = {
        "user": ("users", "username", "User"),
        "account": ("accounts", "account", "Account"),
    }
    resource_mapping = {
        "cpu": "total_consumed_cpu_minutes",
        "gpu": "total_consumed_gpu_minutes",
    }
    if target_type not in target_mapping:
        raise ValueError(f"Invalid target type: {target_type}")
    if resource not in resource_mapping:
        raise ValueError(f"Invalid resource type: {resource}")

    table_name, id_column, target_label = target_mapping[target_type]
    consumed_column = resource_mapping[resource]

    with sqlite3.connect(DB_PATH) as conn:
        cursor = conn.cursor()
        cursor.execute(
            f"SELECT {consumed_column} FROM {table_name} WHERE {id_column} = ?",
            (target_name,),
        )
        row = cursor.fetchone()
        if row is None:
            raise ValueError(f"{target_label} not found: {target_name}")

        current_value = int(row[0] or 0)
        updated_value = max(0, current_value + delta_minutes)
        cursor.execute(
            (
                f"UPDATE {table_name} "
                f"SET {consumed_column} = ?, last_updated = CURRENT_TIMESTAMP "
                f"WHERE {id_column} = ?"
            ),
            (updated_value, target_name),
        )
        conn.commit()
        return updated_value


def create_status_bar(used: int, total: int, width: int = 20) -> str:
    """
    Create a visual status bar representing quota usage.

    Args:
        used: Used amount
        total: Total quota amount
        width: Width of the status bar in characters

    Returns:
        String representation of the status bar
    """
    if total <= 0:
        return " " * width

    percentage = min(used / total, 1.0)
    filled = int(percentage * width)
    bar = "█" * filled + "░" * (width - filled)

    # Apply color based on thresholds unless NO_COLOR is set
    if os.environ.get("NO_COLOR") is None:
        # >95% -> red, >80% -> orange/yellow, else green
        if percentage > 0.95:
            color_start = "\033[31m"  # red
        elif percentage > 0.80:
            # Use 256-color orange when supported; many terminals treat as orange
            color_start = "\033[38;5;208m"  # orange
        else:
            color_start = "\033[32m"  # green
        color_end = "\033[0m"
        colored_bar = f"{color_start}{bar}{color_end}"
    else:
        colored_bar = bar

    return f"[{colored_bar}] {percentage:6.1%}"


class StatsHTTPError(Exception):
    """Raised when the stats HTTP endpoint returns a non-success status."""

    def __init__(self, status: int):
        self.status = status
        super().__init__(f"HTTP {status}")


def fetch_stats_from_service(
    selected_username: Optional[str], selected_account: Optional[str], show_all: bool
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
    """
    Request /stats from the slurm-quota HTTP service and return user and account rows.

    Uses environment variable SLURM_QUOTA_URL as the service root, defaulting to
    http://127.0.0.1:9911/.

    Args:
        selected_username: User filter when not show_all (may be None if show_all).
        selected_account: Account filter to apply on account stats.
        show_all: If True, do not add a username query parameter.

    Returns:
        (users_data, accounts_data) lists from the JSON payload.

    Raises:
        StatsHTTPError: Response status was not HTTP OK.
        URLError: Network or URL handling failure from urlopen.
    """
    base_url = os.environ.get("SLURM_QUOTA_URL", "http://127.0.0.1:9911/")
    stats_params: Dict[str, str] = {}
    if selected_username and not show_all:
        stats_params["username"] = selected_username
    if selected_account:
        stats_params["account"] = selected_account
    stats_url = urljoin(base_url, "stats")
    if stats_params:
        stats_url = f"{stats_url}?{urlencode(stats_params)}"

    with urlopen(Request(stats_url, headers={"Accept": "application/json"})) as resp:
        if resp.status != HTTPStatus.OK:
            raise StatsHTTPError(resp.status)
        stats_payload: Dict[str, Any] = json.load(resp)

    users_data: List[Dict[str, Any]] = list(stats_payload.get("users", []))
    accounts_data: List[Dict[str, Any]] = list(stats_payload.get("accounts", []))
    return users_data, accounts_data


def show_user_stats(
    username: Optional[str] = None,
    account: Optional[str] = None,
    show_all: bool = False,
    display_hours: bool = False,
) -> None:
    """
    Display resource statistics for users in table format.

    Args:
        username: The username to query (if None and not show_all, uses current user)
        account: The account to query in the accounts table.
        show_all: If True, show all users
    """
    # Determine target username
    selected_username = username
    if not show_all and not selected_username and not account:
        try:
            selected_username = get_current_user()
        except KeyError:
            sys.exit(1)

    def format_value(value_minutes: int) -> str:
        if display_hours:
            return f"{value_minutes / 60:.2f}"
        return f"{value_minutes}"

    def compute_first_column_width(
        users: List[Dict[str, Any]],
        accounts: List[Dict[str, Any]],
        max_width: int = 30,
    ) -> int:
        labels = [
            "USER",
            "ACCOUNT",
            *[str(item.get("username", "?")) for item in users],
            *[str(item.get("account", "?")) for item in accounts],
        ]
        longest_label = max(len(label) for label in labels)
        return min(longest_label, max_width)

    try:
        users_data, accounts_data = fetch_stats_from_service(
            selected_username, account, show_all
        )

        if not users_data and not accounts_data:
            if username:
                print(f"No data found for user: {username}")
            elif account:
                print(f"No data found for account: {account}")
            else:
                print("No users found in service")
            return

        first_column_width = compute_first_column_width(users_data, accounts_data)

        def format_first_column(value: str) -> str:
            value_str = str(value)
            if len(value_str) > first_column_width:
                if first_column_width <= 1:
                    value_str = "…" * first_column_width
                else:
                    value_str = f"{value_str[: first_column_width - 1]}…"
            return f"{value_str:<{first_column_width}}"

        # Print users table
        if users_data:
            header = (
                f"{'':<{first_column_width}} | {'CPU':^66} | {'GPU':^66} |\n"
                f"{format_first_column('USER')} | {'CONSUMED':>11} {'PREALLOC(JOBS)':>15} {'QUOTA':>8} {'STATUS':<29} "
                f"| {'CONSUMED':>11} {'PREALLOC(JOBS)':>15} {'QUOTA':>8} {'STATUS':<29} | {'LAST UPDATED':<25}"
            )
            print(header)
            print("-" * int(len(header) / 2))

            for item in users_data:
                uname = item.get("username", "?")
                job_count = int(item.get("job_count", 0))
                last_updated = item.get("last_updated")
                last_updated_str = format_timestamp_with_timezone(last_updated)

                # CPU data
                cpu_consumed = int(item.get("total_consumed_cpu_minutes", 0))
                cpu_preallocated = int(item.get("total_preallocated_cpu_minutes", 0))
                cpu_quota = int(item.get("quota_cpu_minutes", -1))
                cpu_total_used = cpu_consumed + cpu_preallocated
                if cpu_quota == -1:
                    cpu_quota_str = "∞"
                    cpu_status_bar = ""
                else:
                    cpu_quota_str = format_value(cpu_quota)
                    cpu_status_bar = create_status_bar(cpu_total_used, cpu_quota)
                cpu_consumed_str = format_value(cpu_consumed)
                cpu_preallocated_str = f"{format_value(cpu_preallocated)}({job_count})"

                # GPU data
                gpu_consumed = int(item.get("total_consumed_gpu_minutes", 0))
                gpu_preallocated = int(item.get("total_preallocated_gpu_minutes", 0))
                gpu_quota = int(item.get("quota_gpu_minutes", -1))
                gpu_total_used = gpu_consumed + gpu_preallocated
                if gpu_quota == -1:
                    gpu_quota_str = "∞"
                    gpu_status_bar = ""
                else:
                    gpu_quota_str = format_value(gpu_quota)
                    gpu_status_bar = create_status_bar(gpu_total_used, gpu_quota)
                gpu_consumed_str = format_value(gpu_consumed)
                gpu_preallocated_str = f"{format_value(gpu_preallocated)}({job_count})"

                print(
                    f"{format_first_column(uname)} | {cpu_consumed_str:>11} {cpu_preallocated_str:>15} {cpu_quota_str:>8} {cpu_status_bar:<29} "
                    f"| {gpu_consumed_str:>11} {gpu_preallocated_str:>15} {gpu_quota_str:>8} {gpu_status_bar:<29} | {last_updated_str:<25}"
                )
            print()

        # Accounts table
        header = (
            f"{'':<{first_column_width}} | {'CPU':^66} | {'GPU':^66} |\n"
            f"{format_first_column('ACCOUNT')} | {'CONSUMED':>11} {'PREALLOC(JOBS)':>15} {'QUOTA':>8} {'STATUS':<29} "
            f"| {'CONSUMED':>11} {'PREALLOC(JOBS)':>15} {'QUOTA':>8} {'STATUS':<29} | {'LAST UPDATED':<25}"
        )
        print(header)
        print("-" * int(len(header) / 2))

        for item in accounts_data:
            account = item.get("account", "?")
            job_count = int(item.get("job_count", 0))
            last_updated = item.get("last_updated")
            last_updated_str = format_timestamp_with_timezone(last_updated)

            # CPU data
            cpu_consumed = int(item.get("total_consumed_cpu_minutes", 0))
            cpu_preallocated = int(item.get("total_preallocated_cpu_minutes", 0))
            cpu_quota = int(item.get("quota_cpu_minutes", -1))
            cpu_total_used = cpu_consumed + cpu_preallocated
            if cpu_quota == -1:
                cpu_quota_str = "∞"
                cpu_status_bar = " " * 25
            else:
                cpu_quota_str = format_value(cpu_quota)
                cpu_status_bar = create_status_bar(cpu_total_used, cpu_quota)
            cpu_consumed_str = format_value(cpu_consumed)
            cpu_preallocated_str = f"{format_value(cpu_preallocated)}({job_count})"

            # GPU data
            gpu_consumed = int(item.get("total_consumed_gpu_minutes", 0))
            gpu_preallocated = int(item.get("total_preallocated_gpu_minutes", 0))
            gpu_quota = int(item.get("quota_gpu_minutes", -1))
            gpu_total_used = gpu_consumed + gpu_preallocated
            if gpu_quota == -1:
                gpu_quota_str = "∞"
                gpu_status_bar = " " * 25
            else:
                gpu_quota_str = format_value(gpu_quota)
                gpu_status_bar = create_status_bar(gpu_total_used, gpu_quota)
            gpu_consumed_str = format_value(gpu_consumed)
            gpu_preallocated_str = f"{format_value(gpu_preallocated)}({job_count})"

            print(
                f"{format_first_column(account)} | {cpu_consumed_str:>11} {cpu_preallocated_str:>15} {cpu_quota_str:>8} {cpu_status_bar:<29} "
                f"| {gpu_consumed_str:>11} {gpu_preallocated_str:>15} {gpu_quota_str:>8} {gpu_status_bar:<29} | {last_updated_str:<25}"
            )
    except StatsHTTPError as e:
        logger.error(f"Failed to fetch stats: HTTP {e.status}")
        sys.exit(1)
    except URLError as e:
        logger.error(f"Failed to contact slurm-quota service: {e}")
        sys.exit(1)
    except Exception as e:
        logger.error(f"Failed to retrieve stats from service: {e}")
        sys.exit(1)


def query_users_aggregate(
    conn: sqlite3.Connection, username: Optional[str] = None
) -> List[Dict[str, Any]]:
    cursor = conn.cursor()
    if username:
        cursor.execute(
            """
            SELECT u.username, u.total_consumed_cpu_minutes,
                   COALESCE(SUM(j.preallocated_cpu_minutes * COALESCE(j.array_size, 1)), 0) as total_preallocated_cpu_minutes,
                   u.quota_cpu_minutes, u.total_consumed_gpu_minutes,
                   COALESCE(SUM(j.preallocated_gpu_minutes * COALESCE(j.array_size, 1)), 0) as total_preallocated_gpu_minutes,
                   u.quota_gpu_minutes, u.last_updated,
                   COALESCE(SUM(COALESCE(j.array_size, 0)), 0) as job_count
            FROM users u
            LEFT JOIN jobs_preallocations j ON u.username = j.username
            WHERE u.username = ?
            GROUP BY u.username, u.total_consumed_cpu_minutes, u.quota_cpu_minutes,
                     u.total_consumed_gpu_minutes, u.quota_gpu_minutes, u.last_updated
            ORDER BY u.username
            """,
            (username,),
        )
    else:
        cursor.execute(
            """
            SELECT u.username, u.total_consumed_cpu_minutes,
                   COALESCE(SUM(j.preallocated_cpu_minutes * COALESCE(j.array_size, 1)), 0) as total_preallocated_cpu_minutes,
                   u.quota_cpu_minutes, u.total_consumed_gpu_minutes,
                   COALESCE(SUM(j.preallocated_gpu_minutes * COALESCE(j.array_size, 1)), 0) as total_preallocated_gpu_minutes,
                   u.quota_gpu_minutes, u.last_updated,
                   COALESCE(SUM(COALESCE(j.array_size, 0)), 0) as job_count
            FROM users u
            LEFT JOIN jobs_preallocations j ON u.username = j.username
            GROUP BY u.username, u.total_consumed_cpu_minutes, u.quota_cpu_minutes,
                     u.total_consumed_gpu_minutes, u.quota_gpu_minutes, u.last_updated
            ORDER BY u.username
            """
        )
    rows = cursor.fetchall()
    results: List[Dict[str, Any]] = []
    for row in rows:
        (
            uname,
            consumed_cpu,
            prealloc_cpu,
            quota_cpu,
            consumed_gpu,
            prealloc_gpu,
            quota_gpu,
            last_updated,
            job_count,
        ) = row
        results.append(
            {
                "username": uname,
                "total_consumed_cpu_minutes": int(consumed_cpu or 0),
                "total_preallocated_cpu_minutes": int(prealloc_cpu or 0),
                "quota_cpu_minutes": int(quota_cpu if quota_cpu is not None else -1),
                "total_consumed_gpu_minutes": int(consumed_gpu or 0),
                "total_preallocated_gpu_minutes": int(prealloc_gpu or 0),
                "quota_gpu_minutes": int(quota_gpu if quota_gpu is not None else -1),
                "last_updated": last_updated,
                "job_count": int(job_count or 0),
            }
        )
    return results


def query_accounts_aggregate(
    conn: sqlite3.Connection, accounts_filter: Optional[set[str]] = None
) -> List[Dict[str, Any]]:
    cursor = conn.cursor()
    if accounts_filter is not None:
        if len(accounts_filter) == 0:
            return []
        placeholders = ",".join(["?"] * len(accounts_filter))
        cursor.execute(
            f"""
            SELECT a.account, a.total_consumed_cpu_minutes,
                   COALESCE(SUM(j.preallocated_cpu_minutes * COALESCE(j.array_size, 1)), 0) as total_preallocated_cpu_minutes,
                   a.quota_cpu_minutes, a.total_consumed_gpu_minutes,
                   COALESCE(SUM(j.preallocated_gpu_minutes * COALESCE(j.array_size, 1)), 0) as total_preallocated_gpu_minutes,
                   a.quota_gpu_minutes, a.last_updated,
                   COALESCE(SUM(COALESCE(j.array_size, 0)), 0) as job_count
            FROM accounts a
            LEFT JOIN jobs_preallocations j ON a.account = j.account
            WHERE a.account IN ({placeholders})
            GROUP BY a.account, a.total_consumed_cpu_minutes, a.quota_cpu_minutes,
                     a.total_consumed_gpu_minutes, a.quota_gpu_minutes, a.last_updated
            ORDER BY a.account
            """,
            tuple(accounts_filter),
        )
    else:
        cursor.execute(
            """
            SELECT a.account, a.total_consumed_cpu_minutes,
                   COALESCE(SUM(j.preallocated_cpu_minutes * COALESCE(j.array_size, 1)), 0) as total_preallocated_cpu_minutes,
                   a.quota_cpu_minutes, a.total_consumed_gpu_minutes,
                   COALESCE(SUM(j.preallocated_gpu_minutes * COALESCE(j.array_size, 1)), 0) as total_preallocated_gpu_minutes,
                   a.quota_gpu_minutes, a.last_updated,
                   COALESCE(SUM(COALESCE(j.array_size, 0)), 0) as job_count
            FROM accounts a
            LEFT JOIN jobs_preallocations j ON a.account = j.account
            GROUP BY a.account, a.total_consumed_cpu_minutes, a.quota_cpu_minutes,
                     a.total_consumed_gpu_minutes, a.quota_gpu_minutes, a.last_updated
            ORDER BY a.account
            """
        )
    rows = cursor.fetchall()
    results: List[Dict[str, Any]] = []
    for row in rows:
        (
            account,
            consumed_cpu,
            prealloc_cpu,
            quota_cpu,
            consumed_gpu,
            prealloc_gpu,
            quota_gpu,
            last_updated,
            job_count,
        ) = row
        results.append(
            {
                "account": account,
                "total_consumed_cpu_minutes": int(consumed_cpu or 0),
                "total_preallocated_cpu_minutes": int(prealloc_cpu or 0),
                "quota_cpu_minutes": int(quota_cpu if quota_cpu is not None else -1),
                "total_consumed_gpu_minutes": int(consumed_gpu or 0),
                "total_preallocated_gpu_minutes": int(prealloc_gpu or 0),
                "quota_gpu_minutes": int(quota_gpu if quota_gpu is not None else -1),
                "last_updated": last_updated,
                "job_count": int(job_count or 0),
            }
        )
    return results


class _RequestHandler(BaseHTTPRequestHandler):
    server_version = f"slurm-quota-http/{APP_VERSION}"

    def _send_json(self, payload: Any, status: int = 200) -> None:
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, format: str, *args: Any) -> None:
        logger.info("%s - - %s", self.address_string(), format % args)

    def do_GET(self) -> None:
        # Make sure self.server is InactivityHTTPServer with touch_activity method
        assert isinstance(self.server, InactivityHTTPServer)
        # Mark last activity
        self.server.touch_activity()

        if self.path.startswith("/health"):
            self._send_json({"status": "ok"})
            return

        if self.path.startswith("/stats"):
            # Parse query params.
            username_param: Optional[str] = None
            account_param: Optional[str] = None
            query_params = parse_qs(urlparse(self.path).query, keep_blank_values=True)
            usernames = query_params.get("username", [])
            accounts = query_params.get("account", [])
            if usernames:
                username_param = usernames[0]
            if accounts:
                account_param = accounts[0]
            if username_param and account_param:
                self._send_json(
                    {
                        "error": "bad_request",
                        "message": "username and account are mutually exclusive",
                    },
                    status=HTTPStatus.BAD_REQUEST,
                )
                return
            try:
                if not os.path.exists(DB_PATH):
                    self._send_json({"users": [], "accounts": []}, status=200)
                    return
                with sqlite3.connect(DB_PATH) as conn:
                    users = query_users_aggregate(conn, username_param or None)
                    # When a username is provided, filter accounts to that user's associations
                    accounts_filter: Optional[set[str]] = None
                    if username_param:
                        try:
                            accounts_filter = get_user_accounts(username_param)
                        except Exception:
                            accounts_filter = set()
                    if account_param:
                        accounts_filter = {account_param}
                        # Account selection is account-centric; do not include user rows.
                        users = []
                    accounts = query_accounts_aggregate(conn, accounts_filter)
                self._send_json({"users": users, "accounts": accounts}, status=200)
            except sqlite3.Error as e:
                logger.error("/stats query failed: %s", e)
                self._send_json({"error": "db_error"}, status=500)
            return

        # Default 404
        self._send_json({"error": "not_found"}, status=404)


class InactivityHTTPServer(HTTPServer):
    def __init__(self, server_address, RequestHandlerClass, idle_timeout: int = 600):
        super().__init__(server_address, RequestHandlerClass)
        # Get idle timeout from command line or environment variable. Special value 0
        # disables idle shutdown (infinite timeout).
        self._idle_timeout = max(0, int(idle_timeout))
        self._last_activity = time.monotonic()

    def touch_activity(self) -> None:
        self._last_activity = time.monotonic()

    def serve_until_idle(self) -> None:
        self.timeout = 1.0  # seconds
        while True:
            # handle_request() blocks up to self.timeout
            self.handle_request()
            # If idle timeout is disabled, continue to the next iteration.
            if self._idle_timeout == 0:
                continue
            # If idle timeout is enabled, check if the last activity was too long ago.
            if time.monotonic() - self._last_activity > self._idle_timeout:
                logger.info("Idle timeout reached; exiting")
                break


def _systemd_listen_socket() -> Optional[socket.socket]:
    try:
        listen_pid = int(os.environ.get("LISTEN_PID", "0"))
        listen_fds = int(os.environ.get("LISTEN_FDS", "0"))
    except ValueError:
        return None
    if listen_pid != os.getpid() or listen_fds < 1:
        return None
    fd = 3  # per systemd convention
    try:
        # Try AF_INET first; if fails, fallback to generic
        s = socket.fromfd(fd, socket.AF_INET, socket.SOCK_STREAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        return s
    except OSError:
        try:
            s = socket.fromfd(fd, socket.AF_UNIX, socket.SOCK_STREAM)
            return s
        except OSError:
            return None


def run_serve_command(host: str, port: int, idle_timeout: int) -> None:
    # Check if database file exists or exit with error
    if not os.path.exists(DB_PATH):
        logger.error("Database file not found: %s", DB_PATH)
        sys.exit(1)

    # Prefer systemd-provided socket
    sd_sock = _systemd_listen_socket()
    if sd_sock is not None:
        logger.info("Starting HTTP JSON service on systemd socket")
        httpd = InactivityHTTPServer(("0.0.0.0", 0), _RequestHandler, idle_timeout)
        httpd.socket = sd_sock
        # essential: prevent server from closing sd_sock on shutdown duplication issues
        httpd.server_bind = lambda: None  # ty: ignore[invalid-assignment]
        httpd.server_activate = lambda: None  # ty: ignore[invalid-assignment]
        try:
            httpd.serve_until_idle()
        finally:
            try:
                sd_sock.close()
            except Exception:
                pass
        return

    # Fallback: bind our own TCP socket (useful for manual runs)
    addr = (host, int(port))
    logger.info("Starting HTTP JSON service on %s:%s", host, port)
    httpd = InactivityHTTPServer(addr, _RequestHandler, idle_timeout)
    try:
        httpd.serve_until_idle()
    finally:
        try:
            httpd.server_close()
        except Exception:
            pass


def set_user_quota_command(username: str, quota: int) -> None:
    """
    Execute the set-quota command to set user quota.
    This command can only be executed by root user.
    """
    try:
        # Check if running as root user
        current_user = get_current_user()
        if current_user != "root":
            logger.error(
                f"Set-quota command can only be executed by root user, not by "
                f"{current_user}"
            )
            sys.exit(1)

        # Initialize database only if missing
        init_database()

        # Set user quota
        set_user_quota(username, quota)

        print(f"Successfully set quota for user {username}: {quota} CPU minutes")

    except Exception as e:
        logger.error(f"User-quota command failed: {e}")
        sys.exit(1)


def set_account_quota_command(account: str, quota: int) -> None:
    """
    Execute the account-quota command to set account quota.
    This command can only be executed by root user.
    """
    try:
        # Check if running as root user
        current_user = get_current_user()
        if current_user != "root":
            logger.error(
                f"Account-quota command can only be executed by root user, not by "
                f"{current_user}"
            )
            sys.exit(1)

        # Initialize database only if missing
        init_database()

        # Set account quota
        set_account_quota(account, quota)

        print(f"Successfully set quota for account {account}: {quota} CPU minutes")

    except Exception as e:
        logger.error(f"Account-quota command failed: {e}")
        sys.exit(1)


def set_user_gpu_quota_command(username: str, quota: int) -> None:
    """
    Execute the user-gpu-quota command to set user GPU quota.
    This command can only be executed by root user.
    """
    try:
        # Check if running as root user
        current_user = get_current_user()
        if current_user != "root":
            logger.error(
                f"User-gpu-quota command can only be executed by root user, not by "
                f"{current_user}"
            )
            sys.exit(1)

        # Initialize database only if missing
        init_database()

        # Set user GPU quota
        set_user_gpu_quota(username, quota)

        print(f"Successfully set GPU quota for user {username}: {quota} GPU minutes")

    except Exception as e:
        logger.error(f"User-gpu-quota command failed: {e}")
        sys.exit(1)


def set_account_gpu_quota_command(account: str, quota: int) -> None:
    """
    Execute the account-gpu-quota command to set account GPU quota.
    This command can only be executed by root user.
    """
    try:
        # Check if running as root user
        current_user = get_current_user()
        if current_user != "root":
            logger.error(
                f"Account-gpu-quota command can only be executed by root user, not by "
                f"{current_user}"
            )
            sys.exit(1)

        # Initialize database only if missing
        init_database()

        # Set account GPU quota
        set_account_gpu_quota(account, quota)

        print(f"Successfully set GPU quota for account {account}: {quota} GPU minutes")

    except Exception as e:
        logger.error(f"Account-gpu-quota command failed: {e}")
        sys.exit(1)


def adjust_command(
    username: Optional[str],
    account: Optional[str],
    cpu: bool,
    gpu: bool,
    minutes: Optional[int],
    hours: Optional[int],
) -> None:
    """
    Adjust consumed CPU/GPU time for a user or account.
    This command can only be executed by root user.
    """
    try:
        current_user = get_current_user()
        if current_user != "root":
            logger.error(
                "adjust command can only be executed by root user, not by %s",
                current_user,
            )
            sys.exit(1)

        target_type = "user" if username is not None else "account"
        target_name = username if username is not None else account
        if target_name is None:
            logger.error("adjust command requires --user or --account")
            sys.exit(1)

        resource = "cpu" if cpu else "gpu"
        if minutes is not None:
            delta_minutes = minutes
            unit_label = "minutes"
        else:
            if hours is None:
                logger.error("adjust command requires --minutes or --hours")
                sys.exit(1)
            delta_minutes = hours * 60
            unit_label = "hours"

        init_database()
        new_total = adjust_consumed_minutes(
            target_type, target_name, resource, delta_minutes
        )

        delta_source = minutes if minutes is not None else hours
        target_label = "user" if target_type == "user" else "account"
        print(
            (
                f"Successfully adjusted {resource.upper()} consumed time "
                f"for {target_label} {target_name}: "
                f"{delta_source:+d} {unit_label} (new total: {new_total} minutes)"
            )
        )

    except Exception as e:
        logger.error(f"adjust command failed: {e}")
        sys.exit(1)


def show_gpu_factors_command() -> None:
    """
    Execute the gpu-factors command to display current GPU charging factors.
    """
    try:
        if not os.path.exists(DB_PATH):
            print("Database not found. No GPU factors configured.")
            return

        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT gpu_type, factor FROM gpu_factors ORDER BY gpu_type")
            rows = cursor.fetchall()

            if not rows:
                print("No GPU factors configured. Default factor: 1.0")
                return

            # Separate default factor from type-specific factors
            default_factor = None
            type_factors: List[Tuple[str, float]] = []

            for gpu_type, factor in rows:
                try:
                    factor_f = float(factor)
                    if gpu_type == "default":
                        default_factor = factor_f
                    else:
                        type_factors.append((gpu_type, factor_f))
                except (TypeError, ValueError):
                    logger.warning(f"Invalid factor for GPU type {gpu_type}: {factor}")

            # Display results
            print("GPU Charging Factors")
            print("-" * 40)

            # Show default factor
            if default_factor is not None:
                print(f"{'default':<20} {default_factor:.3f}")
            else:
                print(f"{'default':<20} 1.000 (not configured)")

            # Show type-specific factors
            if type_factors:
                for gpu_type, factor in sorted(type_factors):
                    print(f"{gpu_type:<20} {factor:.3f}")

    except sqlite3.Error as e:
        logger.error(f"Failed to query GPU factors: {e}")
        sys.exit(1)
    except Exception as e:
        logger.error(f"gpu-factors command failed: {e}")
        sys.exit(1)


def set_gpu_factor_command(gpu_type: str, factor: float) -> None:
    """
    Execute the set-gpu-factor command to set GPU charging factors.
    This command can only be executed by root user.
    """
    try:
        # Check if running as root user
        current_user = get_current_user()
        if current_user != "root":
            logger.error(
                f"set-gpu-factor command can only be executed by root user, not by "
                f"{current_user}"
            )
            sys.exit(1)

        if factor <= 0:
            logger.error(f"Factor must be positive, got: {factor}")
            sys.exit(1)

        # Initialize database only if missing
        init_database()

        # Set GPU factor
        set_gpu_factor(gpu_type, factor)

        print(f"Successfully set GPU factor for type {gpu_type}: {factor}")

    except Exception as e:
        logger.error(f"set-gpu-factor command failed: {e}")
        sys.exit(1)


def set_default_quotas_command(
    user_cpu: Optional[int],
    user_gpu: Optional[int],
    account_cpu: Optional[int],
    account_gpu: Optional[int],
) -> None:
    """
    Set default quotas used when new users/accounts are auto-created.
    This command can only be executed by root user.
    """
    try:
        current_user = get_current_user()
        if current_user != "root":
            logger.error(
                "set-default-quotas command can only be executed by root user, "
                "not by %s",
                current_user,
            )
            sys.exit(1)

        if all(
            value is None for value in (user_cpu, user_gpu, account_cpu, account_gpu)
        ):
            logger.error(
                "At least one option must be provided: --user-cpu, --user-gpu, "
                "--account-cpu, --account-gpu"
            )
            sys.exit(1)

        init_database()
        set_default_quota_settings(user_cpu, user_gpu, account_cpu, account_gpu)
        print("Successfully updated default quotas")
    except Exception as e:
        logger.error(f"set-default-quotas command failed: {e}")
        sys.exit(1)


def show_default_quotas_command() -> None:
    """
    Display current default quotas for newly auto-created users/accounts.
    """

    def fmt(m: int) -> str:
        return "∞" if m == -1 else str(m)

    try:
        init_database()
        settings = get_default_quota_settings()
        print("Default Quotas For New Entities")
        print("-" * 40)
        print(
            f"{'user_cpu_minutes':<24} "
            f"{fmt(settings['default_user_quota_cpu_minutes'])}"
        )
        print(
            f"{'user_gpu_minutes':<24} "
            f"{fmt(settings['default_user_quota_gpu_minutes'])}"
        )
        print(
            f"{'account_cpu_minutes':<24} "
            f"{fmt(settings['default_account_quota_cpu_minutes'])}"
        )
        print(
            f"{'account_gpu_minutes':<24} "
            f"{fmt(settings['default_account_quota_gpu_minutes'])}"
        )
    except Exception as e:
        logger.error(f"default-quotas command failed: {e}")
        sys.exit(1)


def prune_command(
    preallocs: bool = False,
    users: bool = False,
    accounts: bool = False,
    all_targets: bool = False,
    dry_run: bool = False,
    user_filter: Optional[str] = None,
    account_filter: Optional[str] = None,
) -> None:
    """
    Execute the prune command to remove selected orphan/unused entries.
    This command can only be executed by root user.
    """
    try:
        current_user = get_current_user()
        if current_user != "root":
            logger.error(
                f"Prune command can only be executed by root user, not by {current_user}"
            )
            sys.exit(1)

        if all_targets or not any((preallocs, users, accounts)):
            targets = {"preallocs", "users", "accounts"}
        else:
            targets: set[str] = set()
            if preallocs:
                targets.add("preallocs")
            if users:
                targets.add("users")
            if accounts:
                targets.add("accounts")
        counts = prune_resources(
            targets,
            dry_run=dry_run,
            user_filter=user_filter,
            account_filter=account_filter,
        )
        selected_counts = {target: counts[target] for target in sorted(targets)}

        if dry_run:
            print(
                "Dry-run: would remove "
                f"{selected_counts.get('preallocs', 0)} orphan preallocation(s), "
                f"{selected_counts.get('users', 0)} user(s), "
                f"{selected_counts.get('accounts', 0)} account(s)"
            )
            return

        if all(count == 0 for count in selected_counts.values()):
            print("Nothing to prune")
            return

        print(
            "Removed "
            f"{selected_counts.get('preallocs', 0)} orphan preallocation(s), "
            f"{selected_counts.get('users', 0)} user(s), "
            f"{selected_counts.get('accounts', 0)} account(s)"
        )
    except Exception as e:
        logger.error(f"Prune command failed: {e}")
        sys.exit(1)


def main():
    """
    Main entry point for the slurm-quota script.
    """
    parser = argparse.ArgumentParser(
        prog="slurm-quota",
        description="Slurm Quota Management Tool",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    # Add global logging options
    log_level_group = parser.add_mutually_exclusive_group()
    log_level_group.add_argument(
        "--debug",
        action="store_true",
        help="Print debug output",
    )
    log_level_group.add_argument(
        "-q",
        "--quiet",
        action="store_true",
        help="Only print errors and warnings",
    )
    parser.add_argument(
        "--version",
        action="version",
        version=f"%(prog)s {APP_VERSION}",
        help="Show program version and exit",
    )

    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # Charge command
    subparsers.add_parser(
        "charge", help="Update user resource consumption (for Slurm job completion)"
    )

    # Stats command
    stats_parser = subparsers.add_parser(
        "stats", help="Show resource statistics for users"
    )
    stats_parser.add_argument(
        "username",
        nargs="?",
        help="Username to query (optional, shows all users if not specified)",
    )
    stats_parser.add_argument(
        "--user",
        dest="user",
        help="Username to query (same behavior as positional username)",
    )
    stats_parser.add_argument(
        "--account",
        dest="account",
        help="Account to query",
    )
    stats_parser.add_argument(
        "--all",
        action="store_true",
        help="Show all users (same as not specifying username)",
    )
    stats_parser.add_argument(
        "--hours",
        action="store_true",
        help="Display stats values in hours instead of minutes",
    )

    # Adjust command
    adjust_parser = subparsers.add_parser(
        "adjust",
        help="Adjust consumed CPU/GPU time for a user/account (root only)",
    )
    adjust_target_group = adjust_parser.add_mutually_exclusive_group(required=True)
    adjust_target_group.add_argument("--user", help="Username to adjust")
    adjust_target_group.add_argument("--account", help="Account to adjust")
    adjust_resource_group = adjust_parser.add_mutually_exclusive_group(required=True)
    adjust_resource_group.add_argument(
        "--cpu", action="store_true", help="Adjust consumed CPU minutes"
    )
    adjust_resource_group.add_argument(
        "--gpu", action="store_true", help="Adjust consumed GPU minutes"
    )
    adjust_delta_group = adjust_parser.add_mutually_exclusive_group(required=True)
    adjust_delta_group.add_argument(
        "--minutes",
        type=parse_signed_int,
        help="Signed delta in minutes (must include + or - sign)",
    )
    adjust_delta_group.add_argument(
        "--hours",
        type=parse_signed_int,
        help="Signed delta in hours (must include + or - sign)",
    )

    # User-quota command
    user_quota_parser = subparsers.add_parser(
        "user-quota", help="Set quota for a user (root only)"
    )
    user_quota_parser.add_argument("username", help="Username to set quota for")
    user_quota_parser.add_argument(
        "quota", type=int, help="Quota in CPU minutes (-1 for unlimited)"
    )

    # Account-quota command
    account_quota_parser = subparsers.add_parser(
        "account-quota", help="Set quota for an account (root only)"
    )
    account_quota_parser.add_argument("account", help="Account to set quota for")
    account_quota_parser.add_argument(
        "quota", type=int, help="Quota in CPU minutes (-1 for unlimited)"
    )

    # User-gpu-quota command
    user_gpu_quota_parser = subparsers.add_parser(
        "user-gpu-quota", help="Set GPU quota for a user (root only)"
    )
    user_gpu_quota_parser.add_argument("username", help="Username to set GPU quota for")
    user_gpu_quota_parser.add_argument(
        "quota", type=int, help="Quota in GPU minutes (-1 for unlimited)"
    )

    # Account-gpu-quota command
    account_gpu_quota_parser = subparsers.add_parser(
        "account-gpu-quota", help="Set GPU quota for an account (root only)"
    )
    account_gpu_quota_parser.add_argument(
        "account", help="Account to set GPU quota for"
    )
    account_gpu_quota_parser.add_argument(
        "quota", type=int, help="Quota in GPU minutes (-1 for unlimited)"
    )

    # GPU factors command (show)
    subparsers.add_parser("gpu-factors", help="Show current GPU charging factors")

    # GPU factor command (set)
    gpu_factor_parser = subparsers.add_parser(
        "set-gpu-factor", help="Set GPU charging factor for a GPU type (root only)"
    )
    gpu_factor_parser.add_argument(
        "gpu_type",
        help="GPU type to set factor for (use 'default' for the default factor)",
    )
    gpu_factor_parser.add_argument(
        "factor",
        type=float,
        help="Charging factor (positive float, e.g. 0.5)",
    )

    # Default quotas commands (show)
    subparsers.add_parser(
        "default-quotas",
        help="Show default quotas used for new users/accounts",
    )

    # Default quotas commands (set)
    set_default_quotas_parser = subparsers.add_parser(
        "set-default-quotas",
        help="Set default quotas used for new users/accounts (root only)",
    )
    set_default_quotas_parser.add_argument(
        "--user-cpu",
        type=int,
        help="Default CPU quota for new users in minutes (-1 for unlimited)",
    )
    set_default_quotas_parser.add_argument(
        "--user-gpu",
        type=int,
        help="Default GPU quota for new users in minutes (-1 for unlimited)",
    )
    set_default_quotas_parser.add_argument(
        "--account-cpu",
        type=int,
        help="Default CPU quota for new accounts in minutes (-1 for unlimited)",
    )
    set_default_quotas_parser.add_argument(
        "--account-gpu",
        type=int,
        help="Default GPU quota for new accounts in minutes (-1 for unlimited)",
    )

    # Prune command
    prune_parser = subparsers.add_parser(
        "prune",
        help=("Prune orphaned data from the database (root only)"),
    )
    prune_target_group = prune_parser.add_mutually_exclusive_group()
    prune_target_group.add_argument(
        "--preallocs",
        action="store_true",
        help="Prune orphan preallocations not tied to any active job",
    )
    prune_target_group.add_argument(
        "--users",
        action="store_true",
        help="Prune users with no resource consumption",
    )
    prune_target_group.add_argument(
        "--accounts",
        action="store_true",
        help="Prune accounts with no resource consumption",
    )
    prune_target_group.add_argument(
        "--all",
        action="store_true",
        help="Prune preallocations, users and accounts (default if no selector)",
    )
    prune_parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show how many entries would be removed without deleting anything",
    )
    prune_parser.add_argument(
        "--user",
        help="Limit users pruning to a specific username",
    )
    prune_parser.add_argument(
        "--account",
        help="Limit accounts pruning to a specific account",
    )

    # Serve command (HTTP JSON API for stats, designed for systemd socket activation)
    serve_parser = subparsers.add_parser(
        "serve",
        help="Serve HTTP JSON API (socket-activated by systemd; falls back to host/port)",
    )
    serve_parser.add_argument(
        "--host",
        default="127.0.0.1",
        help="Host to bind if not socket-activated (default: 127.0.0.1)",
    )
    serve_parser.add_argument(
        "--port",
        type=int,
        default=9911,
        help="Port to bind if not socket-activated (default: 9911)",
    )
    serve_parser.add_argument(
        "--idle-timeout",
        type=int,
        default=600,
        help="Exit after N seconds of inactivity (0 disables idle timeout; default: 600)",
    )

    args = parser.parse_args()

    # Setup logging with selected log verbosity
    setup_logging(debug=args.debug, quiet=args.quiet)

    if args.command == "charge":
        charge_command()
    elif args.command == "stats":
        if args.username and args.user:
            parser.error("stats: positional username and --user are mutually exclusive")
        selected_username = args.user or args.username
        if selected_username and args.account:
            parser.error("stats: user selection and --account are mutually exclusive")
        show_user_stats(selected_username, args.account, args.all, args.hours)
    elif args.command == "adjust":
        adjust_command(
            args.user,
            args.account,
            args.cpu,
            args.gpu,
            args.minutes,
            args.hours,
        )
    elif args.command == "user-quota":
        set_user_quota_command(args.username, args.quota)
    elif args.command == "account-quota":
        set_account_quota_command(args.account, args.quota)
    elif args.command == "user-gpu-quota":
        set_user_gpu_quota_command(args.username, args.quota)
    elif args.command == "account-gpu-quota":
        set_account_gpu_quota_command(args.account, args.quota)
    elif args.command == "gpu-factors":
        show_gpu_factors_command()
    elif args.command == "set-gpu-factor":
        set_gpu_factor_command(args.gpu_type, args.factor)
    elif args.command == "default-quotas":
        show_default_quotas_command()
    elif args.command == "set-default-quotas":
        set_default_quotas_command(
            args.user_cpu, args.user_gpu, args.account_cpu, args.account_gpu
        )
    elif args.command == "prune":
        prune_command(
            preallocs=args.preallocs,
            users=args.users,
            accounts=args.accounts,
            all_targets=args.all,
            dry_run=args.dry_run,
            user_filter=args.user,
            account_filter=args.account,
        )
    elif args.command == "serve":
        run_serve_command(args.host, args.port, args.idle_timeout)
    else:
        parser.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()
