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

[(31, 'click.command', 'click.command', ([], {'name': '"""run"""', 'help': '"""Run any daemons configured on the DagsterInstance."""'}), False, 'import click\n'), (51, 'click.command', 'click.command', ([], {'name': '"""health-check"""', 'help': '"""DEPRECATED, use liveness-check instead"""'}), False, 'import click\n'), (65, 'click.command', 'click.command', ([], {'name': '"""liveness-check"""', 'help': '"""Check for recent heartbeats from the daemon."""'}), False, 'import click\n'), (69, '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."""'}), False, 'import click\n'), (84, 'click.command', 'click.command', ([], {'name': '"""wipe"""', 'help': '"""Wipe all heartbeats from storage."""'}), False, 'import click\n'), (94, 'click.command', 'click.command', ([], {'name': '"""heartbeat"""', 'help': '"""Read and write a heartbeat"""'}), False, 'import click\n'), (103, 'click.command', 'click.command', ([], {'name': '"""heartbeat-dump"""', 'help': '"""Log all heartbeat statuses"""'}), False, 'import click\n'), (113, 'click.group', 'click.group', ([], {'commands': "{'heartbeat': debug_heartbeat_command, 'heartbeat-dump':\n    debug_heartbeat_dump_command}"}), False, 'import click\n'), (25, 'os.getenv', 'os.getenv', (['"""DAGSTER_DAEMON_HEARTBEAT_TOLERANCE"""'], {}), False, 'import os\n'), (56, 'warnings.warn', 'warnings.warn', (['"""health-check is deprecated. Use liveness-check instead."""'], {}), False, 'import warnings\n'), (129, 'click.group', 'click.group', ([], {'commands': 'commands'}), False, 'import click\n'), (130, 'click.version_option', 'click.version_option', ([], {'version': '__version__'}), False, 'import click\n'), (36, 'dagster.utils.interrupts.capture_interrupts', 'capture_interrupts', ([], {}), False, 'from dagster.utils.interrupts import capture_interrupts, raise_interrupts_as\n'), (57, 'dagster.core.instance.DagsterInstance.get', 'DagsterInstance.get', ([], {}), False, 'from dagster.core.instance import DagsterInstance\n'), (76, 'dagster.core.instance.DagsterInstance.get', 'DagsterInstance.get', ([], {}), False, 'from dagster.core.instance import DagsterInstance\n'), (89, 'dagster.core.instance.DagsterInstance.get', 'DagsterInstance.get', ([], {}), False, 'from dagster.core.instance import DagsterInstance\n'), (91, 'click.echo', 'click.echo', (['"""Daemon heartbeats wiped"""'], {}), False, 'import click\n'), (99, 'dagster.core.instance.DagsterInstance.get', 'DagsterInstance.get', ([], {}), False, 'from dagster.core.instance import DagsterInstance\n'), (100, 'dagster.daemon.controller.debug_daemon_heartbeats', 'debug_daemon_heartbeats', (['instance'], {}), 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'), (108, 'dagster.core.instance.DagsterInstance.get', 'DagsterInstance.get', ([], {}), False, 'from dagster.core.instance import DagsterInstance\n'), (37, 'dagster.core.instance.DagsterInstance.get', 'DagsterInstance.get', ([], {}), False, 'from dagster.core.instance import DagsterInstance\n'), (59, 'click.echo', 'click.echo', (['"""Daemon healthy"""'], {}), False, 'import click\n'), (61, 'click.echo', 'click.echo', (['"""Daemon not healthy"""'], {}), False, 'import click\n'), (62, 'sys.exit', 'sys.exit', (['(1)'], {}), False, 'import sys\n'), (78, 'click.echo', 'click.echo', (['"""Daemon live"""'], {}), False, 'import click\n'), (80, 'click.echo', 'click.echo', (['"""Daemon(s) not running"""'], {}), False, 'import click\n'), (81, 'sys.exit', 'sys.exit', (['(1)'], {}), False, 'import sys\n'), (110, 'dagster.daemon.controller.get_daemon_status', 'get_daemon_status', (['instance', 'daemon_type'], {}), 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')]