#!/usr/libexec/platform-python

# Copyright 2020 HTCondor Team, Computer Sciences Department,
# University of Wisconsin-Madison, WI.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function

import argparse
import collections
import getpass
import itertools
import sys
import textwrap
import enum
import datetime
import os
import time
import operator
import re
import shutil
import struct

import htcondor2 as htcondor
import classad2 as classad

from htcondor2._utils.ansi import Color, colorize, is_capable
from htcondor2._utils.keyboard import KBHit

IS_WINDOWS = os.name == "nt"


def parse_args():
    parser = argparse.ArgumentParser(
        prog="condor_watch_q",
        description=textwrap.dedent(
            """
            Track the status of jobs over time without repeatedly querying the condor_schedd.

            If no users, cluster ids, event logs, or batch names are given,
            condor_watch_q will default to tracking all of the current user's jobs.

            By default, condor_watch_q will never exit on its own 
            (unless it encounters an error or it is not tracking any jobs).
            You can tell it to exit when certain conditions are met. For example,
            to exit with status 0 when all of the jobs it is tracking are done
            or with status 1 when any job is held, you could run

                condor_watch_q -exit all,done,0 -exit any,held,1
            """
        ),
        formatter_class=argparse.RawTextHelpFormatter,
        add_help=False,
    )

    # GENERAL OPTIONS

    parser.add_argument(
        "-help",
        action="help",
        default=argparse.SUPPRESS,
        help="Show this help message and exit.",
    )

    parser.add_argument(
        "-debug", action="store_true", help="Turn on HTCondor debug printing."
    )

    # TRACKING OPTIONS

    parser.add_argument(
        "-users", nargs="+", metavar="USER", help="Which users to track."
    )
    parser.add_argument(
        "-clusters", nargs="+", metavar="CLUSTER_ID", help="Which cluster IDs to track."
    )
    parser.add_argument(
        "-larger-than",
        type=int,
        dest='since',
        default=None,
        metavar="CLUSTER_ID",
        help="Track all jobs with cluster IDs greater than or equal to the specific CLUSTER_ID."
    )
    parser.add_argument(
        "-files", nargs="+", metavar="FILE", help="Which event logs to track."
    )
    parser.add_argument(
        "-batches", nargs="+", metavar="BATCH_NAME", help="Which batch names to track."
    )

    parser.add_argument(
        "-collector",
        action="store",
        nargs=1,
        default=None,
        help="Which collector to contact to find the schedd, if needed. Defaults to the local collector.",
    )
    parser.add_argument(
        "-schedd",
        action="store",
        nargs=1,
        default=None,
        help="Which schedd to contact for queries, if needed. Defaults to the local schedd.",
    )

    parser.add_argument(
        "-interval",
        action="store",
        default=HALF_A_SECOND,
        type=float,
        help="Interval in seconds between refreshes. Default: %(default)f"
    )

    # BEHAVIOR OPTIONS

    parser.add_argument(
        "-exit",
        action=ExitConditions,
        metavar="GROUPER,JOB_STATUS[,EXIT_STATUS]",
        help=textwrap.dedent(
            """
            Specify conditions under which condor_watch_q should exit. 

            GROUPER is one of {{{exit_groupers}}}.

            JOB_STATUS is one of {{{job_statuses}}}.

            The "active" status means "in the queue",
            and includes jobs in the idle, running, and held states.

            ``EXIT_STATUS`` may be any valid exit status integer.

            To specify multiple exit conditions, pass this option multiple times.
            condor_watch_q will exit when any of the conditions are satisfied.
            """.format(
                exit_groupers=", ".join(EXIT_GROUPERS.keys()),
                job_statuses=", ".join(EXIT_JOB_STATUS_CHECK.keys()),
            )
        ),
    )
    parser.add_argument(
        "-exit-on-key-press",
        "-no-exit-on-key-press",
        action=NegateAction,
        nargs=0,
        default=sys.stdout.isatty(),
        help="Enable/disable exit by any key press. Enabled by default if connected to a tty",
    )

    # DISPLAY OPTIONS

    parser.add_argument(
        "-groupby",
        action="store",
        default="batch",
        choices=("batch", "log", "cluster"),
        help=textwrap.dedent(
            """
            Select what attribute to group jobs by.

            Note that batch names can only be determined if the tracked jobs were
            found in the queue; if they were not, a default batch name is used.
            """
        ),
    )
    parser.add_argument(
        "-display-exited",
        "-no-display-exited",
        action=NegateAction,
        nargs=0,
        default=True,
        help="Enable/disable displaying rows of jobs that have all exited the queue. Enabled by default.",
    )
    parser.add_argument(
        "-exited-row-decay",
        action=ExitDisplayDecay,
        metavar="TIME[unit]",
        nargs=1,
        default=(7 * 24 * 60 * 60),
        help=textwrap.dedent(
            """
            Time since the last job in a job exited the queue before no
            longer than being displayed (if -no-display-exited specified).
            Default time is one week. Unless specified default unit is in
            weeks.

            Specifiable units:
                s = Seconds
                m = Minutes
                h = Hours
                d = Days
                w = Weeks
                M = Months
                y = years
            """
        )
    )
    parser.add_argument(
        "-decay-start-at-read",
        "-no-decay-start-at-read",
        action=NegateAction,
        nargs=0,
        default=False,
        help="Start all job exited row display decay timer at event read time rather than event write time. Disabled by default",
    )
    parser.add_argument(
        "-table",
        "-no-table",
        action=NegateAction,
        nargs=0,
        default=True,
        help="Enable/disable the table. Enabled by default.",
    )
    parser.add_argument(
        "-progress",
        "-no-progress",
        action=NegateAction,
        nargs=0,
        default=True,
        help="Enable/disable the progress bar. Enabled by default.",
    )
    parser.add_argument(
        "-row-progress",
        "-no-row-progress",
        action=NegateAction,
        nargs=0,
        default=True,
        help="Enable/disable the progress bar for each row. Enabled by default.",
    )
    parser.add_argument(
        "-summary",
        "-no-summary",
        action=NegateAction,
        nargs=0,
        default=True,
        help="Enable/disable the summary line. Enabled by default.",
    )
    parser.add_argument(
        "-summary-type",
        action="store",
        choices=("totals", "percentages"),
        default="totals",
        help="Choose what to display on the summary line. By default, show totals.",
    )
    parser.add_argument(
        "-updated-at",
        "-no-updated-at",
        action=NegateAction,
        nargs=0,
        default=True,
        help='Enable/disable the "updated at" line. Enabled by default.',
    )
    parser.add_argument(
        "-abbreviate",
        "-no-abbreviate",
        action=NegateAction,
        nargs=0,
        default=True,
        help="Enable/disable abbreviating path components to the shortest somewhat-unique prefix. Disabled by default.",
    )
    parser.add_argument(
        "-color",
        "-no-color",
        action=NegateAction,
        nargs=0,
        default=is_capable(),
        help="Enable/disable colored output. Enabled by default if connected to a tty. Disabled on Windows if colorama is not available.",
    )
    parser.add_argument(
        "-refresh",
        "-no-refresh",
        action=NegateAction,
        nargs=0,
        default=sys.stdout.isatty(),
        help="Enable/disable refreshing output (instead of appending). Enabled by default if connected to a tty.",
    )


    args, unknown = parser.parse_known_args()

    if len(unknown) != 0:
        check_unknown_args_for_known_errors(parser, unknown)

        # because we're in this branch, this will fail, printing out the
        # original errors argparse would have printed for the unknown args
        # if we didn't catch them ourselves
        args = parser.parse_args()

    args.groupby = {
        "log": "event_log_path",
        "cluster": "cluster_id",
        "batch": "batch_name",
    }[args.groupby]

    return args


def check_unknown_args_for_known_errors(parser, unknown_args):
    """
    This function tries to capture known errors in the arguments that
    argparse fails to parse and emit our own error messages.
    Individual args are passed to _check_unknown_arg for processing.
    Common errors are trying to use condor_q-like arguments.
    """
    unknown_args = iter(unknown_args)
    for arg in unknown_args:
        error_message = _check_unknown_arg(arg, unknown_args)
        if error_message is None:
            continue

        parser.print_usage()
        print(
            "{}: error: argument {}: {}".format(parser.prog, arg, error_message),
            file=sys.stderr,
        )
        sys.exit(1)


def _check_unknown_arg(arg, unknown_args):
    if arg.isdigit():
        cluster_ids = [arg]
        for next_arg in unknown_args:
            if next_arg.startswith("-"):
                break
            if not next_arg.isdigit():
                continue
            cluster_ids.append(next_arg)

        return "to track specific cluster IDs, try -clusters {}".format(
            " ".join(cluster_ids)
        )

    elif "-totals" in arg:
        return "to only print totals, try -no-table"

    elif "-userlog" in arg:
        files = []
        for next_arg in unknown_args:
            if next_arg.startswith("-"):
                break
            files.append(next_arg)

        return "to track jobs from specific event logs, try -files {}".format(
            " ".join(files)
        )

    elif "-nobatch" in arg:
        return "to group by something other than batch name, try -groupby {{batch,log,cluster}}"

    elif not arg.startswith("-"):
        users = [arg]
        for next_arg in unknown_args:
            if next_arg.startswith("-"):
                break
            if next_arg.isdigit():
                continue
            users.append(next_arg)

        return "to track jobs from specific users, try -users {}".format(
            " ".join(users)
        )


class ExitConditions(argparse.Action):
    def __call__(self, parser, args, values, option_string=None):
        v = values.split(",")
        if len(v) == 3:
            grouper, status, exit_status = v
        elif len(v) == 2:
            grouper, status = v
            exit_status = 0
        else:
            parser.error(message="invalid -exit specification")

        if grouper not in EXIT_GROUPERS:
            parser.error(
                message='invalid GROUPER "{grouper}", must be one of {{{groupers}}}'.format(
                    grouper=grouper, groupers=", ".join(EXIT_GROUPERS.keys())
                )
            )

        if status not in EXIT_JOB_STATUS_CHECK:
            parser.error(
                message='invalid JOB_STATUS "{status}", must be one of {{{statuses}}}'.format(
                    status=status, statuses=", ".join(EXIT_JOB_STATUS_CHECK.keys())
                )
            )

        try:
            exit_status = int(exit_status)
        except ValueError:
            parser.error(
                message='exit_status must be an integer, but was "{}"'.format(
                    exit_status
                )
            )

        if getattr(args, self.dest, None) is None:
            setattr(args, self.dest, [])

        getattr(args, self.dest).append((grouper, status, exit_status))


class NegateAction(argparse.Action):
    def __init__(self, option_strings, dest, **kwargs):
        super().__init__(option_strings, dest, **kwargs)

    def __call__(self, parser, args, values, option_string=None):
        enable = (option_string.strip("-").lower().startswith("no") != True)
        setattr(args, self.dest, enable)


class ExitDisplayDecay(argparse.Action):
    UNIT_CONVERSION = {
        "s": 1,
        "m": 60,
        "h": 60 * 60,
        "d": 24 * 60 * 60,
        "w": 7 * 24 * 60 * 60,
        "M": 31 * 24 * 60 * 60,
        "y": 365 * 24 * 60 * 60,
    }

    def __init__(self, option_strings, dest, nargs=None, **kwargs):
            if nargs is None:
                raise ValueError("Requires one argument")
            elif int(nargs) != 1:
                raise ValueError("Only one argument allowed")
            super().__init__(option_strings, dest, **kwargs)

    def __call__(self, parser, args, values, option_string=None):
        unit = "w"
        if not values[-1].isdigit():
            unit = values[-1]
            values = values[:-1]

        if unit not in self.UNIT_CONVERSION:
            parser.error(message=f"Unknown unit specifier provided for -exited-row-decay: {unit}")
        else:
            try:
                decay = int(values) * self.UNIT_CONVERSION[unit]
            except ValueError:
                parser.error(message=f"Invalid time value provided for -exited-row-decay")

        setattr(args, self.dest, decay)


def cli():
    args = parse_args()

    if args.clusters is not None and args.since is not None:
        error("Use of both -clusters and -larger-than flag is prohibited.")
        sys.exit(1)

    if args.debug:
        warning("Debug mode enabled...")
        htcondor.enable_debug()

    try:
        return watch_q(
            users=args.users,
            cluster_ids=args.clusters,
            since_cid=args.since,
            specified_logs=args.files,
            batches=args.batches,
            collector=args.collector,
            schedd=args.schedd,
            exit_conditions=args.exit,
            group_by=args.groupby,
            table=args.table,
            progress_bar=args.progress,
            row_progress_bar=args.row_progress,
            summary=args.summary,
            summary_type=args.summary_type,
            updated_at=args.updated_at,
            color=args.color,
            refresh=args.refresh,
            abbreviate_path_components=args.abbreviate,
            args=args
        )
    except Exception as e:
        if args.debug:
            raise

        error(
            "Unhandled error: {}. Re-run with -debug for a full stack trace.".format(e)
        )
        sys.exit(1)


EXIT_GROUPERS = {"all": all, "any": any, "none": lambda _: not any(_)}
EXIT_JOB_STATUS_CHECK = {
    "active": lambda s: s in ACTIVE_STATES,
    "done": lambda s: s is JobStatus.COMPLETED,
    "idle": lambda s: s is JobStatus.IDLE,
    "held": lambda s: s is JobStatus.HELD,
}

TOTAL = "TOTAL"
ACTIVE_JOBS = "JOB_IDS"
EVENT_LOG = "LOG"
CLUSTER_ID = "CLUSTER"
BATCH_NAME = "BATCH"

# attribute is the Python attribute name of the Cluster object
# key is the key in the events and job rows
GROUPBY_ATTRIBUTE_TO_AD_KEY = {
    "event_log_path": EVENT_LOG,
    "cluster_id": CLUSTER_ID,
    "batch_name": BATCH_NAME,
}
GROUPBY_AD_KEY_TO_ATTRIBUTE = {v: k for k, v in GROUPBY_ATTRIBUTE_TO_AD_KEY.items()}

HALF_A_SECOND = 500 / 1000

# Escape sequences to switch to/from the "alternate screen" buffer, the same
# mechanism `less` and other full-screen programs use so that the terminal
# returns to its prior contents once the program exits.
ALT_SCREEN_ENTER = "\033[?1049h\033[H"
ALT_SCREEN_EXIT = "\033[?1049l"


def watch_q(
    users=None,
    cluster_ids=None,
    since_cid=None,
    specified_logs=None,
    batches=None,
    collector=None,
    schedd=None,
    exit_conditions=None,
    group_by="batch_name",
    table=True,
    progress_bar=True,
    row_progress_bar=True,
    summary=True,
    summary_type="totals",
    updated_at=True,
    color=True,
    refresh=True,
    abbreviate_path_components=False,
    args=None,
):
    filter_ids = None if cluster_ids is None else [int(cid) for cid in cluster_ids]
    if users is None and cluster_ids is None and specified_logs is None and batches is None:
        users = [getpass.getuser()]
    if exit_conditions is None:
        exit_conditions = []

    key = GROUPBY_ATTRIBUTE_TO_AD_KEY[group_by]

    row_fmt = (lambda s, r: colorize(s, determine_row_color(r))) if color else None

    specified_clusters = (cluster_ids is not None)

    (
        cluster_ids,
        num_procs,
        constraint,
        event_logs,
        batch_names,
        batch_filter_ids,
        dagman_clusters_to_paths,
        dagman_job_cluster_ids,
    ) = find_job_event_logs(
        users, cluster_ids, specified_logs, batches, collector=collector, schedd=schedd,
    )

    if filter_ids is None:
        filter_ids = batch_filter_ids if len(batch_filter_ids) > 0 else None
    else:
        filter_ids.extend(batch_filter_ids)

    tracker = JobStateTracker(event_logs, batch_names, filter_ids, since_cid, num_procs, args.decay_start_at_read)
    if len(event_logs) == 0:
        warning("No jobs found, exiting...")
        sys.exit(0)

    exit_checks = []
    for grouper, checker, exit_status in exit_conditions:
        disp = "{} {}".format(grouper, checker)
        exit_grouper = EXIT_GROUPERS[grouper.lower()]
        exit_check = EXIT_JOB_STATUS_CHECK[checker.lower()]
        exit_checks.append((exit_grouper, exit_check, exit_status, disp))

    # When refreshing in place (i.e. on a tty), draw into the terminal's
    # "alternate screen" buffer so that quitting restores whatever was on the
    # console before condor_watch_q started, the way `less` behaves.
    refresh = refresh and sys.stdout.isatty()

    in_alt_screen = False

    def leave_alt_screen():
        nonlocal in_alt_screen
        if in_alt_screen:
            sys.stdout.write(ALT_SCREEN_EXIT)
            sys.stdout.flush()
            in_alt_screen = False

    if refresh:
        sys.stdout.write(ALT_SCREEN_ENTER)
        sys.stdout.flush()
        in_alt_screen = True

        import signal

        def _handle_exit_signal(signum, frame):
            leave_alt_screen()
            raise KeyboardInterrupt

        for _sig in ("SIGTERM", "SIGHUP", "SIGQUIT"):
            if hasattr(signal, _sig):
                signal.signal(getattr(signal, _sig), _handle_exit_signal)
    try:
        msg, move, clear = None, None, None
        batch_name_to_nodes_total, cluster_ids_with_nodes_total = find_nodes_total(
            constraint, collector, schedd
        )
        last_nodes_total_check_time = time.time()

        keyboard = KBHit(args.exit_on_key_press)

        while True:
            processing_messages = tracker.process_events()
            now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

            if (
                not all(
                    cluster_id in cluster_ids_with_nodes_total
                    for cluster_id in dagman_job_cluster_ids
                )
                and time.time() - last_nodes_total_check_time > 60
            ):
                last_nodes_total_check_time = time.time()
                (
                    batch_name_to_nodes_total,
                    cluster_ids_with_nodes_total,
                ) = find_nodes_total(constraint, collector, schedd)

            # "if msg is not None" skips the first iteration, when there's nothing to clear
            if msg is not None and refresh:
                msg = strip_ansi(msg)
                prev_lines = list(msg.splitlines())
                prev_len_lines = [len(line) for line in prev_lines]
                move = "\033[{}A\r".format(len(prev_len_lines))
                clear = "\n".join(" " * len(line) for line in prev_lines) + "\n"

            groups, batch_name_to_nodes_total = group_clusters(
                tracker.clusters,
                key,
                dagman_clusters_to_paths,
                batch_names,
                dagman_job_cluster_ids,
                batch_name_to_nodes_total,
                specified_clusters,
            )
            rows_by_key, totals = make_rows_from_groups(
                groups, key, batch_name_to_nodes_total, args
            )

            headers, rows_by_key = strip_empty_columns(rows_by_key)

            # strip out 0 values
            rows_by_key = {
                key: {k: v for k, v in row.items() if v != 0}
                for key, row in rows_by_key.items()
            }

            if key == EVENT_LOG and abbreviate_path_components:
                for row in rows_by_key.values():
                    row[key] = abbreviate_path(row[key])

            # Sort so active groups print first
            rows_by_key = sorted(
                rows_by_key.items(),
                key= lambda item: row_sorter(item, groups),
            )

            msg = []

            terminal_columns, terminal_rows = shutil.get_terminal_size()

            if table:
                msg += make_table(
                    row_progress_bar=row_progress_bar,
                    terminal_columns=terminal_columns,
                    headers=[key] + headers,
                    rows=[row for _, row in rows_by_key],
                    row_fmt=row_fmt,
                    alignment=TABLE_ALIGNMENT,
                    fill="-",
                    color=color,
                )
                msg += [""]

            if progress_bar:
                msg += make_total_progress_bar(totals, terminal_columns, color)
                msg += [""]

            if summary:
                if summary_type == "totals":
                    msg += make_summary_with_totals(totals, width=terminal_columns)
                elif summary_type == "percentages":
                    msg += make_summary_with_percentages(
                        totals, width=terminal_columns
                    )
                msg += [""]

            if updated_at:
                msg += ["Updated at {}".format(now)] + [""]

            if len(msg) > terminal_rows:
                msg = msg[: terminal_rows - 2]
                msg = (
                    "\n".join(msg[:-1])
                    + "\nInsufficient terminal height to display full output!"
                )
            else:
                # msg[:-1] because we need to strip the last blank section delimiter line off
                msg = "\n".join(msg[:-1])

            if not refresh:
                msg += "\n..."

            if not exit_conditions:
                msg += "\nPress any key to exit" if args.exit_on_key_press else "\nPress ^C to exit"

            if clear and move:
                sys.stdout.write(move + clear + move)
                sys.stdout.flush()

            if len(processing_messages) > 0:
                print(
                    "\n".join("{} | {}".format(now, m) for m in processing_messages),
                    file=sys.stderr,
                )

            print(msg)

            for grouper, checker, exit_status, disp in exit_checks:
                if grouper((checker(s) for s in tracker.job_states)):
                    # Restore the original screen before printing the reason so
                    # it remains visible after we exit.
                    leave_alt_screen()
                    print(
                        'Exiting with code {} because of condition "{}" at {}'.format(
                            exit_status, disp, now
                        )
                    )
                    sys.exit(exit_status)

            if args.exit_on_key_press and keyboard.anyKeyPressed():
                raise KeyboardInterrupt

            time.sleep(args.interval)
    except KeyboardInterrupt:
        pass
    finally:
        leave_alt_screen()
    sys.exit(0)


PROJECTION = [
    "ClusterId",
    "Owner",
    "UserLog",
    "DAGManNodesLog",
    "JobBatchName",
    "Iwd",
    "DAG_NodesTotal",
    "OtherJobRemoveRequirements",
    "TotalSubmitProcs",
]


def find_job_event_logs(
    users=None, cluster_ids=None, files=None, batches=None, collector=None, schedd=None,
):
    """
    Discover job event logs to read events from based on various methods.

    Parameters
    ----------
    users
        Find job event logs for these user's active jobs.
    cluster_ids
        Find job event logs for these clusters.
    files
        Find these job event logs (basically, these just get passed straight through).
    batches
        Find job event logs for these batch names.
    collector
        Query this collector to find the schedd.
        Defaults to the local collector.
    schedd
        Query this schedd for users, cluster_ids, and batches.
        Defaults to the local schedd.
    """
    if users is None:
        users = []
    if cluster_ids is None:
        cluster_ids = []
    if files is None:
        files = []
    if batches is None:
        batches = []

    constraint = " || ".join(
        itertools.chain(
            ("Owner == {}".format(classad.quote(u)) for u in users),
            ("ClusterId == {}".format(cid) for cid in cluster_ids),
            ('JobBatchName == "{}"'.format(b) for b in batches),
        )
    )

    clusters = set()
    num_procs = dict()
    event_logs = dict()
    batch_names = {}
    batch_filter_ids = []
    already_warned_missing_log = set()
    dagman_job_cluster_id_to_log_path = {}
    dagman_job_cluster_ids = set()

    for file in files:
        event_logs[os.path.abspath(file)] = False

    for ad in get_ads(constraint, collector, schedd):
        cluster_id = ad["ClusterId"]
        clusters.add(cluster_id)
        num_procs[cluster_id] = ad["TotalSubmitProcs"]

        # Don't do ad.get() because we call str() on the return
        batch_names[cluster_id] = str(ad["JobBatchName"]) if "JobBatchName" in ad else None

        if len(batches) > 0:
            batch_filter_ids.append(int(cluster_id))

        if "DAGManNodesLog" in ad:
            log_path = dagman_job_cluster_id_to_log_path[cluster_id] = ad[
                "DAGManNodesLog"
            ]
        elif "UserLog" in ad:
            log_path = ad["UserLog"]
        else:
            if cluster_id not in already_warned_missing_log:
                warning(
                    "Cluster {} does not have a job event log file (set log=<path> in the submit description)".format(
                        cluster_id
                    )
                )
                already_warned_missing_log.add(cluster_id)
            continue

        if not os.path.isabs(log_path):
            log_path = os.path.abspath(os.path.join(ad["Iwd"], log_path))

        event_logs[log_path] = False

        # this job is the actual DAGMan controller job
        if "OtherJobRemoveRequirements" in ad:
            dagman_job_cluster_ids.add(cluster_id)
            nodes_log = log_path.replace(".dagman.log", ".nodes.log")
            if nodes_log not in event_logs:
                event_logs[nodes_log] = True

    return (
        clusters,
        num_procs,
        constraint,
        event_logs,
        batch_names,
        batch_filter_ids,
        dagman_job_cluster_id_to_log_path,
        dagman_job_cluster_ids,
    )


def find_nodes_total(constraint, collector, schedd):
    batch_name_to_nodes_total = {}
    cluster_ids_with_nodes_total = set()

    for ad in get_ads(constraint, collector, schedd):
        try:
            job_batch_name = ad["JobBatchName"]
            cluster_id = ad["ClusterId"]
        except KeyError:
            continue

        if "DAG_NodesTotal" in ad and ad["DAG_NodesTotal"]:
            batch_name_to_nodes_total[job_batch_name] = ad["DAG_NodesTotal"]
            cluster_ids_with_nodes_total.add(cluster_id)

    return batch_name_to_nodes_total, cluster_ids_with_nodes_total


def get_ads(constraint, collector, schedd):
    if constraint == "":
        return []

    try:
        return get_schedd(collector=collector, schedd=schedd).query(
            constraint, PROJECTION
        )
    except Exception as e:
        warning("Was not able to query to discover job event logs due to: {}".format(e))
        return []


def get_schedd(collector=None, schedd=None):
    if collector is None and schedd is None:
        schedd = htcondor.Schedd()
    else:
        coll = htcondor.Collector(collector)
        schedd_ad = coll.locate(htcondor.DaemonTypes.Schedd, schedd)
        schedd = htcondor.Schedd(schedd_ad)

    return schedd


class Cluster:
    """Holds the job state for a singe cluster."""

    def __init__(self, cluster_id, event_log_path, batch_name, proc_count):
        self.cluster_id = cluster_id
        self.event_log_path = event_log_path
        self._batch_name = batch_name
        self._proc_count = proc_count

        self.job_to_state = {}
        self.decay_start = None

    @property
    def num_jobs(self):
        return self._proc_count if self._proc_count is not None else len(self.job_to_state)

    @property
    def batch_name(self):
        return self._batch_name or "ID: {}".format(self.cluster_id)

    def __setitem__(self, key, value):
        self.job_to_state[key] = value

    def __getitem__(self, item):
        return self.job_to_state[item]

    def items(self):
        return self.job_to_state.items()

    def __iter__(self):
        return iter(self.items())

    def __lt__(self, other):
        return self.cluster_id < other.cluster_id

    def done(self) -> bool:
        """Cluster is considered done if all associated jobs have left the queue (completed or removed state)"""
        return len(self.job_to_state) == sum(
            [
                1 if state in EXITED_STATES else 0 for state in self.job_to_state.values()
            ]
        )

    def display(self, decay_time) -> bool:
        """Determine whether cluster should be display based on jobs exiting queue and decay times"""
        if self.done() and self.decay_start is not None:
            return time.time() - self.decay_start < decay_time

        return True


class JobStateTracker:
    """
    Tracks the job state from many event logs,
    maintaining a mapping of cluster_id to Cluster
    (available as the state attribute).
    """

    def __init__(self, event_log_paths, batch_names, filter_ids, since_cid, num_procs, decay_start):
        self.event_readers = dict()
        self.deferred_logs = self.open_readers(event_log_paths)
        self.state = collections.defaultdict(lambda: collections.defaultdict(dict))

        self.batch_names = batch_names
        self.filter_ids = filter_ids
        self.since = since_cid

        self.cluster_id_to_cluster = {}
        self.num_procs_by_cluster = num_procs
        self.decay_start_at_read = decay_start

    def open_readers(self, logs: dict) -> list:
        deferred_logs = dict()
        for log, defer in logs.items():
            try:
                reader = htcondor.JobEventLog(log).events(0)
                self.event_readers[log] = reader
            except Exception as e:
                if defer and log.endswith(".nodes.log"):
                    deferred_logs[log] = defer
                    continue
                if not os.path.exists(log):
                    e = "File path does not exist"
                warning(
                    f"Could not open event log at {log} for reading, so it will be ignored. Reason: {e}"
                )
        return deferred_logs

    def process_events(self):
        """
        Process all currently-available events in all event logs.
        """
        messages = []

        for event_log_path, events in self.event_readers.items():
            while True:
                try:
                    event = next(events)
                except StopIteration:
                    break
                except Exception as e:
                    messages.append(
                        fmt_error(
                            "Failed to parse event from {}. Reason: {}".format(
                                event_log_path, e
                            )
                        )
                    )
                    continue

                if self.filter_ids is not None and event.cluster not in self.filter_ids:
                    continue

                if self.since is not None and event.cluster < self.since:
                    continue

                if event.type == htcondor.JobEventType.FILE_TRANSFER:
                    new_status = None
                    msg = str(event).lower()
                    if "started" in msg:
                        if "input" in msg:
                            new_status = JobStatus.TRANSFERRING_INPUT
                        elif "output" in msg:
                            new_status = JobStatus.TRANSFERRING_OUTPUT
                else:
                    new_status = JOB_EVENT_STATUS_TRANSITIONS.get(event.type, None)

                if new_status is None:
                    continue

                # Check submit type events for structured notes for a batch name if missing
                batch_name = self.batch_names.get(event.cluster)
                if batch_name is None and event.type in [htcondor.JobEventType.SUBMIT, htcondor.JobEventType.CLUSTER_SUBMIT]:
                    note_attrs = event.get("StructuredNotes")
                    if note_attrs is not None:
                        batch_name = note_attrs.get("JobBatchName")
                        self.batch_names[event.cluster] = batch_name

                cluster = self.cluster_id_to_cluster.setdefault(
                    event.cluster,
                    Cluster(
                        cluster_id=event.cluster,
                        event_log_path=event_log_path,
                        batch_name=batch_name,
                        proc_count=self.num_procs_by_cluster.get(event.cluster),
                    ),
                )

                cluster[event.proc] = new_status

                if new_status in EXITED_STATES and cluster.done():
                    cluster.decay_start = time.time() if self.decay_start_at_read else event.timestamp

        if len(self.deferred_logs) > 0:
            self.deferred_logs = self.open_readers(self.deferred_logs)

        return messages

    @property
    def clusters(self):
        return self.cluster_id_to_cluster.values()

    @property
    def job_states(self):
        for cluster in self.clusters:
            for job, state in cluster:
                yield state


def group_clusters(
    clusters,
    key,
    dagman_job_cluster_id_to_log_path,
    batch_names,
    dagman_job_cluster_ids,
    dag_nodes_total,
    specified_clusters,
):
    groups = collections.defaultdict(list)
    getter = operator.attrgetter(GROUPBY_AD_KEY_TO_ATTRIBUTE[key])
    group_nodes_total = {}

    dag_batch_names_by_nodes_log = {
        nodes_log_path: batch_names[cluster]
        for cluster, nodes_log_path in dagman_job_cluster_id_to_log_path.items()
    }
    for cluster in clusters:
        # find cluster according to cluster path
        dag_batch_name = dag_batch_names_by_nodes_log.get(cluster.event_log_path)
        if key == BATCH_NAME and dag_batch_name is not None:
            groups[dag_batch_name].append(cluster)
            group_nodes_total[dag_batch_name] = dag_nodes_total.get(dag_batch_name)
        elif specified_clusters or cluster.cluster_id not in dagman_job_cluster_ids:
            groups[getter(cluster)].append(cluster)

    if not dag_nodes_total:
        group_nodes_total = dag_nodes_total

    return groups, group_nodes_total


def make_rows_from_groups(groups, key, cluster_id_to_nodes_total, args):
    totals = collections.defaultdict(int)
    rows = {}

    for attribute_value, clusters in groups.items():
        row_data = row_data_from_job_state(
            clusters, cluster_id_to_nodes_total.get(attribute_value)
        )

        totals[TOTAL] += row_data[TOTAL]
        for status in JobStatus:
            totals[status] += row_data[status]

        if key == EVENT_LOG:
            row_data[key] = nice_path(attribute_value)
        else:
            row_data[key] = attribute_value

        # Filter out done groups
        if args.display_exited or any([c.display(args.exited_row_decay) for c in clusters]):
            rows[attribute_value] = row_data

    return rows, totals


def row_sorter(item, groups) -> int:
    clusters = groups[item[0]]
    # Find the cluster with the smallest ClusterId for ordering
    canidate = min(clusters)

    # If active jobs then make sort key value a negative number
    # relative to cluster id such that smaller ids are sorted first
    if any([not c.done() for c in clusters]):
        return (sys.maxsize - 1 - canidate.cluster_id) * -1

    return canidate.cluster_id


def determine_row_color(row):
    if row.get(JobStatus.HELD, 0) > 0:
        return Color.RED
    elif row.get(JobStatus.COMPLETED) == row.get("TOTAL"):
        return Color.GREEN
    elif row.get(JobStatus.RUNNING, 0) > 0:
        return Color.CYAN
    elif row.get(JobStatus.IDLE, 0) > 0:
        return Color.YELLOW
    else:
        return Color.BRIGHT_WHITE


def strip_empty_columns(rows_by_key):
    dont_include = set()
    for h in HEADERS:
        if all((row[h] == 0 for row in rows_by_key.values())):
            dont_include.add(h)
    dont_include -= ALWAYS_INCLUDE
    headers = [h for h in HEADERS if h not in dont_include]
    for row_data in dont_include:
        for row in rows_by_key.values():
            row.pop(row_data)

    return headers, rows_by_key


class JobId():
    def __init__(self, cluster: int, proc: int) -> None:
        self.cluster = cluster
        self.proc = proc

    def __str__(self) -> str:
        return f"{self.cluster}.{self.proc}"

    def __lt__(self, other) -> bool:
        return (self.cluster, self.proc) < (other.cluster, other.proc)


def row_data_from_job_state(clusters, expected_total_jobs=None):
    """Construct the data for a single row."""
    if expected_total_jobs is None:
        # ensure expected_total_jobs will get clobbered by the max below
        expected_total_jobs = -1

    row_data = {js: 0 for js in JobStatus}
    active_job_ids = []

    total_jobs = 0
    for cluster in clusters:
        total_jobs += cluster.num_jobs
        for proc_id, job_state in cluster:
            row_data[job_state] += 1

            if job_state in ACTIVE_STATES:
                active_job_ids.append(JobId(cluster.cluster_id, proc_id))

    row_data[TOTAL] = max(expected_total_jobs, total_jobs)
    active_job_ids.sort()

    if len(active_job_ids) > 2:
        active_job_ids = [active_job_ids[0], active_job_ids[-1]]
        row_data[ACTIVE_JOBS] = " ... ".join(str(id) for id in active_job_ids)
    else:
        row_data[ACTIVE_JOBS] = ", ".join(str(id) for id in active_job_ids)

    return row_data


def nice_path(path):
    """
    Try to find a "nice" representation of a path.
    The shortest (by string length) of these options is returned:
      - the absolute path
      - the path relative to the current user's home directory
      - the path relative to the current working directory
    """
    possibilities = []

    abs = os.path.abspath(path)
    home_dir = os.path.expanduser("~")
    cwd = os.getcwd()

    possibilities.append(abs)

    relative_to_user_home = os.path.relpath(path, home_dir)
    possibilities.append(os.path.join("~", relative_to_user_home))

    relative_to_cwd = os.path.relpath(path, cwd)
    if not relative_to_cwd.startswith("../"):  # i.e, it is not *above* the cwd
        possibilities.append(os.path.join(".", relative_to_cwd))

    return min(possibilities, key=len)


def abbreviate_path(path):
    """
    Abbreviate a filesystem path by finding somewhat-unique prefixes for each
    component in the path.

    If the path was /home/foobar/wizbang/events.log and there was nothing else in any
    of those directories, the result would be /h/f/w/events.log . If there was a
    directory /home/foobar/wizard, the result would be /h/f/wizb/events.log .
    """
    abbreviated_components = []
    path_so_far = ""
    components = split_all(path)
    for component in components[:-1]:
        # the actual, unabbreviated path to this depth in the path
        path_so_far = os.path.expanduser(os.path.join(path_so_far, component))

        # these have special meaning, so don't do anything to them
        if component in ("~", "."):
            abbreviated_components.append(component)
            continue

        contents = os.listdir(os.path.dirname(path_so_far))
        if len(contents) == 1:
            longest_common = ""
        else:
            longest_common = os.path.commonprefix(contents)

        # tack on the common prefix plus one extra character to give a degree of uniqueness
        abbreviated_components.append(component[: len(longest_common) + 1])

    # tack on the last path component as-is; this is the actual event log file name
    abbreviated_components.append(components[-1])

    return os.path.join(*abbreviated_components)


def split_all(path):
    """Consistently split a path into components."""
    split_parts = []
    while 1:
        parts = os.path.split(path)
        if parts[0] == path:  # sentinel for absolute paths
            split_parts.insert(0, parts[0])
            break
        elif parts[1] == path:  # sentinel for relative paths
            split_parts.insert(0, parts[1])
            break
        else:
            path = parts[0]
            split_parts.insert(0, parts[1])
    return split_parts


class JobStatus(enum.Enum):
    REMOVED = "REMOVED"
    HELD = "HELD"
    IDLE = "IDLE"
    RUNNING = "RUN"
    COMPLETED = "DONE"
    TRANSFERRING_INPUT = "INPUT"
    TRANSFERRING_OUTPUT = "OUTPUT"
    SUSPENDED = "SUSPENDED"

    def __str__(self):
        return self.value

    @classmethod
    def ordered(cls):
        return (
            cls.REMOVED,
            cls.HELD,
            cls.SUSPENDED,
            cls.IDLE,
            cls.TRANSFERRING_INPUT,
            cls.RUNNING,
            cls.TRANSFERRING_OUTPUT,
            cls.COMPLETED,
        )


ACTIVE_STATES = {
    JobStatus.IDLE,
    JobStatus.TRANSFERRING_INPUT,
    JobStatus.RUNNING,
    JobStatus.TRANSFERRING_OUTPUT,
    JobStatus.HELD,
    JobStatus.SUSPENDED,
}

EXITED_STATES = {
    JobStatus.REMOVED,
    JobStatus.COMPLETED,
}

ALWAYS_INCLUDE = {
    JobStatus.IDLE,
    JobStatus.RUNNING,
    JobStatus.COMPLETED,
    EVENT_LOG,
    TOTAL,
}

TABLE_ALIGNMENT = {
    EVENT_LOG: "ljust",
    CLUSTER_ID: "ljust",
    TOTAL: "rjust",
    ACTIVE_JOBS: "ljust",
    BATCH_NAME: "ljust",
}
for k in JobStatus:
    TABLE_ALIGNMENT[k] = "rjust"

HEADERS = list(JobStatus.ordered()) + [TOTAL, ACTIVE_JOBS]

# NOTE: TRANSFERRING_(IN/OUT)PUT are handle specially in process_events()
#       due to all transfers sharing a single job event type (#40)
JOB_EVENT_STATUS_TRANSITIONS = {
    htcondor.JobEventType.SUBMIT: JobStatus.IDLE,
    htcondor.JobEventType.JOB_EVICTED: JobStatus.IDLE,
    htcondor.JobEventType.JOB_UNSUSPENDED: JobStatus.IDLE,
    htcondor.JobEventType.JOB_RELEASED: JobStatus.IDLE,
    htcondor.JobEventType.SHADOW_EXCEPTION: JobStatus.IDLE,
    htcondor.JobEventType.JOB_RECONNECT_FAILED: JobStatus.IDLE,
    htcondor.JobEventType.JOB_TERMINATED: JobStatus.COMPLETED,
    htcondor.JobEventType.EXECUTE: JobStatus.RUNNING,
    htcondor.JobEventType.JOB_HELD: JobStatus.HELD,
    htcondor.JobEventType.JOB_SUSPENDED: JobStatus.SUSPENDED,
    htcondor.JobEventType.JOB_ABORTED: JobStatus.REMOVED,
}


def make_table(
    row_progress_bar,
    terminal_columns,
    headers,
    rows,
    fill="",
    header_fmt=None,
    row_fmt=None,
    alignment=None,
    color=True,
):
    if header_fmt is None:
        header_fmt = lambda _: _
    if row_fmt is None:
        row_fmt = lambda _a, _b: _a
    if alignment is None:
        alignment = {}

    headers = tuple(headers)
    lengths = [len(str(h)) for h in headers]
    align_methods = [alignment.get(h, "center") for h in headers]
    processed_rows = []

    for row in rows:
        processed_rows.append([str(row.get(key, fill)) + " " for key in headers])
    for row in processed_rows:
        lengths = [max(curr, len(entry)) for curr, entry in zip(lengths, row)]

    header = header_fmt(
        "  ".join(
            getattr(str(h), a)(l) for h, l, a in zip(headers, lengths, align_methods)
        ).rstrip()
    )

    lines = [
        row_fmt(
            "  ".join(
                getattr(f, a)(l)
                for f, l, a in zip(processed_row, lengths, align_methods)
            ),
            original_row,
        )
        for original_row, processed_row in zip(rows, processed_rows)
    ]

    try:
        remaining_columns = terminal_columns - len(strip_ansi(lines[0]))
    except IndexError:
        remaining_columns = terminal_columns

    if row_progress_bar and remaining_columns > 10:
        for i, row in zip(range(len(lines)), rows):
            lines[i] += "".join(make_progress_bar(row, remaining_columns, color=color))
    else:
        # Iterate through every row in table, truncate to console width
        for row in range(len(lines)):
            lines[row] = lines[row][:terminal_columns]

    return [header] + lines


ANSI_ESCAPE_RE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")


def strip_ansi(string):
    return ANSI_ESCAPE_RE.sub("", string)


PROGRESS_BAR_CHARS_AND_COLORS = {
    JobStatus.IDLE: (Color.BRIGHT_YELLOW, "-"),
    JobStatus.TRANSFERRING_INPUT: (Color.CYAN, ">"),
    JobStatus.RUNNING: (Color.CYAN, "="),
    JobStatus.TRANSFERRING_OUTPUT: (Color.CYAN, ">"),
    JobStatus.COMPLETED: (Color.GREEN, "#"),
    JobStatus.HELD: (Color.RED, "!"),
    JobStatus.SUSPENDED: (Color.RED, "!"),
    JobStatus.REMOVED: (Color.RED, "!"),
}

PROGRESS_BAR_ORDER = (
    JobStatus.IDLE,
    JobStatus.TRANSFERRING_INPUT,
    JobStatus.RUNNING,
    JobStatus.TRANSFERRING_OUTPUT,
    JobStatus.COMPLETED,
    JobStatus.HELD,
    JobStatus.SUSPENDED,
    JobStatus.REMOVED,
)


def make_progress_bar(totals, width=None, color=True):
    width = min(width, 79) - 2  # account for the wrapping [ ]
    num_total = float(totals[TOTAL])

    fractions = [
        safe_divide(totals.get(status), num_total)
        for status in PROGRESS_BAR_ORDER
    ]

    bar_section_lengths = [int(width * f) for f in fractions]

    # give any rounded-off space to the longest part of the bar
    bar_section_lengths[argmax(bar_section_lengths)] += width - sum(bar_section_lengths)

    bar = "[{}]".format(
        "".join(
            colorize(
                PROGRESS_BAR_CHARS_AND_COLORS[status][1] * length,
                PROGRESS_BAR_CHARS_AND_COLORS[status][0],
            )
            if color
            else PROGRESS_BAR_CHARS_AND_COLORS[status][1] * length
            for status, length in zip(PROGRESS_BAR_ORDER, bar_section_lengths)
        )
    )

    return [bar]

def make_total_progress_bar(totals, width=None, color=True):
    if width is not None:
        width = width - 16

    result = "Total Progress: " + make_progress_bar(totals, width, color)[0]
    return [result]


def argmax(collection):
    """
    Return the index of the element of the collection with the maximum value.
    """
    return max(enumerate(collection), key=lambda x: x[1])[0]


# this controls the ordering of the summary message
SUMMARY_STATES = [
    JobStatus.COMPLETED,
    JobStatus.REMOVED,
    JobStatus.IDLE,
    # Running jobs include jobs doing transfers
    [
        JobStatus.TRANSFERRING_INPUT,
        JobStatus.RUNNING,
        JobStatus.TRANSFERRING_OUTPUT,
    ],
    JobStatus.HELD,
    JobStatus.SUSPENDED,
]
SUMMARY_MESSAGES = ["completed", "removed", "idle", "running", "held", "suspended"]

# Calculate summary states totals: Either 1 state or tuple of states
def calc_multi_state_total(totals, states) -> int:
    if isinstance(states, list):
        return sum([totals[status] for status in states])
    elif isinstance(states, JobStatus):
        return totals[states]
    else:
        # Developer error
        raise RuntimeError("Invalid states variable type (expected JobStatus or tuple[JobStatus])")


def make_summary_with_totals(totals, width=None):
    strip = 11

    while True:
        strip -= 1

        summary = "Total: {} jobs; {}".format(
            totals[TOTAL],
            ", ".join(
                "{} {}".format(calc_multi_state_total(totals, status), msg[:strip])
                for status, msg in zip(SUMMARY_STATES, SUMMARY_MESSAGES)
                if calc_multi_state_total(totals, status) > 0
            ),
        )

        if width is None or len(summary) <= width or strip <= 2:
            break

    return [summary]


def make_summary_with_percentages(totals, width=None):
    num_total = totals[TOTAL]
    percentages = [
        round(100.0 * safe_divide(calc_multi_state_total(totals, s), num_total)) for s in SUMMARY_STATES
    ]

    summary = "Total: {} jobs; {}".format(
        totals[TOTAL],
        ", ".join(
            "{}% {}".format(p, msg) for p, msg in zip(percentages, SUMMARY_MESSAGES)
        ),
    )

    if width is not None and len(summary[0]) > width:
        summary = "Total: {} jobs; {}".format(
            totals[TOTAL],
            ", ".join(
                "{}% {}".format(p, msg[0].upper())
                for p, msg in zip(percentages, SUMMARY_MESSAGES)
            ),
        )

    return [summary]


def safe_divide(numerator, denominator, default=0):
    if not numerator or not denominator:
        return default
    try:
        return numerator / denominator
    except ZeroDivisionError:
        return default


def warning(msg):
    print(fmt_warning(msg), file=sys.stderr)


def fmt_warning(msg):
    return "WARNING: {}".format(msg)


def error(msg):
    print(fmt_error(msg), file=sys.stderr)


def fmt_error(msg):
    return "ERROR: {}".format(msg)


if __name__ == "__main__":
    cli()
