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', False, 'import click\n'), (51, 'click.command', 'click.command', False, 'import click\n'), (65, 'click.command', 'click.command', False, 'import click\n'), (69, 'click.option', 'click.option', False, 'import click\n'), (84, 'click.command', 'click.command', False, 'import click\n'), (94, 'click.command', 'click.command', False, 'import click\n'), (103, 'click.command', 'click.command', False, 'import click\n'), (113, 'click.group', 'click.group', False, 'import click\n'), (25, 'os.getenv', 'os.getenv', False, 'import os\n'), (56, 'warnings.warn', 'warnings.warn', False, 'import warnings\n'), (129, 'click.group', 'click.group', False, 'import click\n'), (130, 'click.version_option', 'click.version_option', 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', 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', 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', False, 'import click\n'), (61, 'click.echo', 'click.echo', False, 'import click\n'), (62, 'sys.exit', 'sys.exit', False, 'import sys\n'), (78, 'click.echo', 'click.echo', False, 'import click\n'), (80, 'click.echo', 'click.echo', False, 'import click\n'), (81, 'sys.exit', 'sys.exit', False, 'import sys\n'), (110, 'dagster.daemon.controller.get_daemon_status', 'get_daemon_status', 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')]