import os
import sys
import threading
import time
import warnings
from contextlib import ExitStack

import click
import pendulum
from dagster import __version__
from dagster.core.instance import DagsterInstance
from dagster.daemon.controller import (
    DEFAULT_DAEMON_HEARTBEAT_TOLERANCE_SECONDS,
    DagsterDaemonController,
    all_daemons_healthy,
    all_daemons_live,
    daemon_controller_from_instance,
    debug_daemon_heartbeats,
    get_daemon_status,
)
from dagster.utils.interrupts import capture_interrupts, raise_interrupts_as


def _get_heartbeat_tolerance():
    tolerance = os.getenv(
        "DAGSTER_DAEMON_HEARTBEAT_TOLERANCE",
    )
    return int(tolerance) if tolerance else DEFAULT_DAEMON_HEARTBEAT_TOLERANCE_SECONDS


@click.command(
    name="run",
    help="Run any daemons configured on the DagsterInstance.",
)
def run_command():
    with capture_interrupts():
        with DagsterInstance.get() as instance:
            if instance.is_ephemeral:
                raise Exception(
                    "dagster-daemon can't run using an in-memory instance. Make sure "
                    "the DAGSTER_HOME environment variable has been set correctly and that "
                    "you have created a dagster.yaml file there."
                )

            with daemon_controller_from_instance(
                instance, heartbeat_tolerance_seconds=_get_heartbeat_tolerance()
            ) as controller:
                controller.check_daemon_loop()


@click.command(
    name="health-check",
    help="DEPRECATED, use liveness-check instead",
)
def health_check_command():
    warnings.warn("health-check is deprecated. Use liveness-check instead.")
    with DagsterInstance.get() as instance:
        if all_daemons_healthy(instance, heartbeat_tolerance_seconds=_get_heartbeat_tolerance()):
            click.echo("Daemon healthy")
        else:
            click.echo("Daemon not healthy")
            sys.exit(1)


@click.command(
    name="liveness-check",
    help="Check for recent heartbeats from the daemon.",
)
@click.option(
    "--heartbeat-tolerance",
    required=False,
    default=DEFAULT_DAEMON_HEARTBEAT_TOLERANCE_SECONDS,
    help="How long (in seconds) to allow a daemon to go without heartbeating before failing the dagster-daemon process.",
)
def liveness_check_command():
    with DagsterInstance.get() as instance:
        if all_daemons_live(instance, heartbeat_tolerance_seconds=_get_heartbeat_tolerance()):
            click.echo("Daemon live")
        else:
            click.echo("Daemon(s) not running")
            sys.exit(1)


@click.command(
    name="wipe",
    help="Wipe all heartbeats from storage.",
)
def wipe_command():
    with DagsterInstance.get() as instance:
        instance.wipe_daemon_heartbeats()
        click.echo("Daemon heartbeats wiped")


@click.command(
    name="heartbeat",
    help="Read and write a heartbeat",
)
def debug_heartbeat_command():
    with DagsterInstance.get() as instance:
        debug_daemon_heartbeats(instance)


@click.command(
    name="heartbeat-dump",
    help="Log all heartbeat statuses",
)
def debug_heartbeat_dump_command():
    with DagsterInstance.get() as instance:
        for daemon_type in instance.get_required_daemon_types():
            click.echo(get_daemon_status(instance, daemon_type))


@click.group(
    commands={"heartbeat": debug_heartbeat_command, "heartbeat-dump": debug_heartbeat_dump_command}
)
def debug_group():
    "Daemon debugging utils"


def create_dagster_daemon_cli():
    commands = {
        "run": run_command,
        "health-check": health_check_command,
        "liveness-check": liveness_check_command,
        "wipe": wipe_command,
        "debug": debug_group,
    }

    @click.group(commands=commands)
    @click.version_option(version=__version__)
    def group():
        "CLI tools for working with the dagster daemon process."

    return group


cli = create_dagster_daemon_cli()


def main():
    cli(obj={})  # pylint:disable=E1123

[((746, 835), 'click.command', 'click.command', ([], {'name': '"""run"""', 'help': '"""Run any daemons configured on the DagsterInstance."""'}), "(name='run', help=\n    'Run any daemons configured on the DagsterInstance.')\n", (759, 835), False, 'import click\n'), ((1486, 1572), 'click.command', 'click.command', ([], {'name': '"""health-check"""', 'help': '"""DEPRECATED, use liveness-check instead"""'}), "(name='health-check', help=\n    'DEPRECATED, use liveness-check instead')\n", (1499, 1572), False, 'import click\n'), ((1953, 2047), 'click.command', 'click.command', ([], {'name': '"""liveness-check"""', 'help': '"""Check for recent heartbeats from the daemon."""'}), "(name='liveness-check', help=\n    'Check for recent heartbeats from the daemon.')\n", (1966, 2047), False, 'import click\n'), ((2055, 2293), 'click.option', 'click.option', (['"""--heartbeat-tolerance"""'], {'required': '(False)', 'default': 'DEFAULT_DAEMON_HEARTBEAT_TOLERANCE_SECONDS', 'help': '"""How long (in seconds) to allow a daemon to go without heartbeating before failing the dagster-daemon process."""'}), "('--heartbeat-tolerance', required=False, default=\n    DEFAULT_DAEMON_HEARTBEAT_TOLERANCE_SECONDS, help=\n    'How long (in seconds) to allow a daemon to go without heartbeating before failing the dagster-daemon process.'\n    )\n", (2067, 2293), False, 'import click\n'), ((2594, 2662), 'click.command', 'click.command', ([], {'name': '"""wipe"""', 'help': '"""Wipe all heartbeats from storage."""'}), "(name='wipe', help='Wipe all heartbeats from storage.')\n", (2607, 2662), False, 'import click\n'), ((2829, 2895), 'click.command', 'click.command', ([], {'name': '"""heartbeat"""', 'help': '"""Read and write a heartbeat"""'}), "(name='heartbeat', help='Read and write a heartbeat')\n", (2842, 2895), False, 'import click\n'), ((3027, 3098), 'click.command', 'click.command', ([], {'name': '"""heartbeat-dump"""', 'help': '"""Log all heartbeat statuses"""'}), "(name='heartbeat-dump', help='Log all heartbeat statuses')\n", (3040, 3098), False, 'import click\n'), ((3323, 3435), 'click.group', 'click.group', ([], {'commands': "{'heartbeat': debug_heartbeat_command, 'heartbeat-dump':\n    debug_heartbeat_dump_command}"}), "(commands={'heartbeat': debug_heartbeat_command,\n    'heartbeat-dump': debug_heartbeat_dump_command})\n", (3334, 3435), False, 'import click\n'), ((593, 640), 'os.getenv', 'os.getenv', (['"""DAGSTER_DAEMON_HEARTBEAT_TOLERANCE"""'], {}), "('DAGSTER_DAEMON_HEARTBEAT_TOLERANCE')\n", (602, 640), False, 'import os\n'), ((1611, 1683), 'warnings.warn', 'warnings.warn', (['"""health-check is deprecated. Use liveness-check instead."""'], {}), "('health-check is deprecated. Use liveness-check instead.')\n", (1624, 1683), False, 'import warnings\n'), ((3734, 3764), 'click.group', 'click.group', ([], {'commands': 'commands'}), '(commands=commands)\n', (3745, 3764), False, 'import click\n'), ((3770, 3811), 'click.version_option', 'click.version_option', ([], {'version': '__version__'}), '(version=__version__)\n', (3790, 3811), False, 'import click\n'), ((870, 890), 'dagster.utils.interrupts.capture_interrupts', 'capture_interrupts', ([], {}), '()\n', (888, 890), False, 'from dagster.utils.interrupts import capture_interrupts, raise_interrupts_as\n'), ((1693, 1714), 'dagster.core.instance.DagsterInstance.get', 'DagsterInstance.get', ([], {}), '()\n', (1712, 1714), False, 'from dagster.core.instance import DagsterInstance\n'), ((2337, 2358), 'dagster.core.instance.DagsterInstance.get', 'DagsterInstance.get', ([], {}), '()\n', (2356, 2358), False, 'from dagster.core.instance import DagsterInstance\n'), ((2703, 2724), 'dagster.core.instance.DagsterInstance.get', 'DagsterInstance.get', ([], {}), '()\n', (2722, 2724), False, 'from dagster.core.instance import DagsterInstance\n'), ((2788, 2825), 'click.echo', 'click.echo', (['"""Daemon heartbeats wiped"""'], {}), "('Daemon heartbeats wiped')\n", (2798, 2825), False, 'import click\n'), ((2947, 2968), 'dagster.core.instance.DagsterInstance.get', 'DagsterInstance.get', ([], {}), '()\n', (2966, 2968), False, 'from dagster.core.instance import DagsterInstance\n'), ((2990, 3023), 'dagster.daemon.controller.debug_daemon_heartbeats', 'debug_daemon_heartbeats', (['instance'], {}), '(instance)\n', (3013, 3023), False, 'from dagster.daemon.controller import DEFAULT_DAEMON_HEARTBEAT_TOLERANCE_SECONDS, DagsterDaemonController, all_daemons_healthy, all_daemons_live, daemon_controller_from_instance, debug_daemon_heartbeats, get_daemon_status\n'), ((3155, 3176), 'dagster.core.instance.DagsterInstance.get', 'DagsterInstance.get', ([], {}), '()\n', (3174, 3176), False, 'from dagster.core.instance import DagsterInstance\n'), ((905, 926), 'dagster.core.instance.DagsterInstance.get', 'DagsterInstance.get', ([], {}), '()\n', (924, 926), False, 'from dagster.core.instance import DagsterInstance\n'), ((1838, 1866), 'click.echo', 'click.echo', (['"""Daemon healthy"""'], {}), "('Daemon healthy')\n", (1848, 1866), False, 'import click\n'), ((1893, 1925), 'click.echo', 'click.echo', (['"""Daemon not healthy"""'], {}), "('Daemon not healthy')\n", (1903, 1925), False, 'import click\n'), ((1938, 1949), 'sys.exit', 'sys.exit', (['(1)'], {}), '(1)\n', (1946, 1949), False, 'import sys\n'), ((2479, 2504), 'click.echo', 'click.echo', (['"""Daemon live"""'], {}), "('Daemon live')\n", (2489, 2504), False, 'import click\n'), ((2531, 2566), 'click.echo', 'click.echo', (['"""Daemon(s) not running"""'], {}), "('Daemon(s) not running')\n", (2541, 2566), False, 'import click\n'), ((2579, 2590), 'sys.exit', 'sys.exit', (['(1)'], {}), '(1)\n', (2587, 2590), False, 'import sys\n'), ((3278, 3318), 'dagster.daemon.controller.get_daemon_status', 'get_daemon_status', (['instance', 'daemon_type'], {}), '(instance, daemon_type)\n', (3295, 3318), False, 'from dagster.daemon.controller import DEFAULT_DAEMON_HEARTBEAT_TOLERANCE_SECONDS, DagsterDaemonController, all_daemons_healthy, all_daemons_live, daemon_controller_from_instance, debug_daemon_heartbeats, get_daemon_status\n')]