#  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
#
#  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:
#
#       https://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.
"""ZenML Step HeartBeat functionality."""

import _thread
import contextvars
import logging
import threading
import time
from datetime import datetime, timezone
from typing import TYPE_CHECKING
from uuid import UUID

from zenml.enums import ExecutionStatus
from zenml.utils.time_utils import to_utc_timezone

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)


class StepHeartBeatTerminationException(Exception):
    """Custom exception class for heartbeat termination."""

    pass


class StepHeartbeatWorker:
    """Worker class implementing heartbeat polling and remote termination."""

    STEP_HEARTBEAT_INTERVAL_SECONDS = (
        60 * 2
    )  # submit heartbeats every 2 minutes.

    MAX_HEARTBEAT_INTERVAL_SECONDS = 60 * 5
    MAX_CONSECUTIVE_FAILURES = 3

    def __init__(self, step_id: UUID):
        """Heartbeat worker constructor.

        Args:
            step_id: The step id heartbeat is running for.
        """
        self._step_id = step_id

        self._thread: threading.Thread | None = None
        self._running: bool = False
        self._terminated: bool = False

        self._consecutive_failures: int = 0
        self._heartbeat_interval_seconds: int = (
            self.STEP_HEARTBEAT_INTERVAL_SECONDS
        )

    # properties

    @property
    def is_running(self) -> bool:
        """Property function for running signal.

        Returns:
            True if the worker is running.
        """
        return self._running

    @property
    def is_terminated(self) -> bool:
        """Property function for termination signal.

        Returns:
            True if the worker has been terminated.
        """
        return self._terminated

    @property
    def interval(self) -> int:
        """Property function for heartbeat interval.

        Returns:
            The heartbeat polling interval value.
        """
        return self._heartbeat_interval_seconds

    @property
    def name(self) -> str:
        """Property function for heartbeat worker name.

        Returns:
            The name of the heartbeat worker.
        """
        return f"HeartBeatWorker-{self.step_id}"

    @property
    def step_id(self) -> UUID:
        """Property function for heartbeat worker step ID.

        Returns:
            The id of the step heartbeat is running for.
        """
        return self._step_id

    # public functions

    def start(self) -> None:
        """Start the heartbeat worker on a background thread."""
        if self._thread and self._thread.is_alive():
            return

        self._running = True
        self._terminated = False

        active_context = contextvars.copy_context()
        self._thread = threading.Thread(
            target=lambda: active_context.run(self._run),
            name=self.name,
            daemon=True,
        )
        self._thread.start()
        logger.debug(
            "Daemon thread %s started (interval=%s)", self.name, self.interval
        )

    def stop(self) -> None:
        """Stops the heartbeat worker."""
        if not self._running:
            return
        self._running = False
        logger.debug("%s stop requested", self.name)

    def is_alive(self) -> bool:
        """Liveness of the heartbeat worker thread.

        Returns:
            True if the heartbeat worker thread is alive, False otherwise.
        """
        t = self._thread
        return bool(t and t.is_alive())

    def _run(self) -> None:
        logger.debug("%s run() loop entered", self.name)
        try:
            while self._running:
                try:
                    self._heartbeat()
                    self._consecutive_failures = 0
                    # One-shot: signal the main thread and stop the loop.
                    if self._terminated:
                        logger.info(
                            "%s is remotely stopped, interrupting main thread",
                            self.name,
                        )
                        _thread.interrupt_main()  # raises KeyboardInterrupt in main thread
                        self._running = False
                    # Ensure we stop our own loop as well.

                except Exception as exc:
                    self._consecutive_failures += 1
                    # Log-and-continue policy for all other errors.
                    logger.debug(
                        "%s heartbeat() failed with %s", self.name, str(exc)
                    )

                    if (
                        self._consecutive_failures
                        == self.MAX_CONSECUTIVE_FAILURES
                    ):
                        logger.debug(
                            "Reached %s consecutive failures - applying increased heartbeat interval..",
                            self.MAX_CONSECUTIVE_FAILURES,
                        )
                        self._heartbeat_interval_seconds = (
                            self.MAX_HEARTBEAT_INTERVAL_SECONDS
                        )

                # Sleep after each attempt (even after errors, unless stopped).
                if self._running:
                    time.sleep(self.interval)
        finally:
            logger.debug("%s run() loop exiting", self.name)

    def _heartbeat(self) -> None:
        from zenml.config.global_config import GlobalConfiguration

        store = GlobalConfiguration().zen_store

        response = store.update_step_heartbeat(step_run_id=self.step_id)

        if not response.heartbeat_enabled:
            logger.debug("Heartbeat set to disabled - stopping worker")
            self.stop()
            return

        if response.pipeline_run_status in {
            ExecutionStatus.STOPPED,
            ExecutionStatus.STOPPING,
        }:
            self._terminated = True


def is_heartbeat_unhealthy(
    step_run_id: UUID,
    status: ExecutionStatus,
    latest_heartbeat: datetime | None,
    start_time: datetime | None = None,
    heartbeat_threshold: int | None = None,
) -> bool:
    """Utility function - Checks if step heartbeats indicate un-healthy execution.

    Args:
        step_run_id: The run step id.
        status: The run step status.
        latest_heartbeat: The run step latest heartbeat.
        start_time: The run step start time.
        heartbeat_threshold: If heartbeat enabled the max minutes without heartbeat
            for healthy, running tasks.

    Returns:
        True if the step heartbeat is unhealthy, False otherwise.
    """
    if not heartbeat_threshold:
        return False

    if status.is_finished:
        return False

    if latest_heartbeat:
        heartbeat_diff = datetime.now(tz=timezone.utc) - to_utc_timezone(
            latest_heartbeat
        )
    elif start_time:
        heartbeat_diff = datetime.now(tz=timezone.utc) - to_utc_timezone(
            start_time
        )
    else:
        return False

    logger.debug("Step %s heartbeat diff=%s", step_run_id, heartbeat_diff)

    if heartbeat_diff.total_seconds() > heartbeat_threshold * 60:
        return True

    return False
