"""
This module defines the Heartbeat class.
The heartbeat is responsible for sending periodic pings to the Runpod server.
"""

import os
import time
from multiprocessing import Process
import requests
from urllib3.util.retry import Retry

from runpod.http_client import SyncClientSession
from runpod.serverless.modules.rp_logger import RunPodLogger
from runpod.serverless.modules.worker_state import WORKER_ID
from runpod.version import __version__ as runpod_version

log = RunPodLogger()


class Heartbeat:
    """Sends heartbeats to the Runpod server."""

    _process_started = False

    def __init__(self, pool_connections=10, retries=3) -> None:
        """
        Initializes the Heartbeat class.
        """
        self.PING_URL = os.environ.get("RUNPOD_WEBHOOK_PING", "PING_NOT_SET")
        self.PING_URL = self.PING_URL.replace("$RUNPOD_POD_ID", WORKER_ID)
        self.PING_INTERVAL = int(os.environ.get("RUNPOD_PING_INTERVAL", 10000)) // 1000

        # In-progress job-id snapshot, injected by the main process at start.
        self._mirror = None

        # Create a new HTTP session
        self._session = SyncClientSession()
        self._session.headers.update(
            {"Authorization": os.environ.get("RUNPOD_AI_API_KEY", "")}
        )

        retry_strategy = Retry(
            total=retries,
            status_forcelist=[429, 500, 502, 503, 504],
            allowed_methods=["GET"],
            backoff_factor=1,
        )

        adapter = requests.adapters.HTTPAdapter(
            pool_connections=pool_connections,
            pool_maxsize=pool_connections,
            max_retries=retry_strategy,
        )
        self._session.mount("http://", adapter)
        self._session.mount("https://", adapter)

    @staticmethod
    def process_loop(mirror=None, test=False):
        """
        Static helper to run the ping loop in a separate process.
        Creates a new Heartbeat instance to avoid pickling issues.
        """
        hb = Heartbeat()
        hb._mirror = mirror
        hb.ping_loop(test)

    def start_ping(self, mirror=None, test=False):
        """
        Sends heartbeat pings to the Runpod server in a separate process.
        """
        if not os.environ.get("RUNPOD_AI_API_KEY"):
            log.debug("Not deployed on Runpod serverless, pings will not be sent.")
            return

        if not os.environ.get("RUNPOD_POD_ID"):
            log.info("Not running on Runpod, pings will not be sent.")
            return

        if (not self.PING_URL) or self.PING_URL == "PING_NOT_SET":
            log.error("Ping URL not set, cannot start ping.")
            return

        if not Heartbeat._process_started:
            process = Process(target=Heartbeat.process_loop, args=(mirror, test))
            process.daemon = True
            process.start()
            Heartbeat._process_started = True

    def ping_loop(self, test=False):
        """
        Sends heartbeat pings to the Runpod server until interrupted.
        """
        while True:
            self._send_ping()
            time.sleep(self.PING_INTERVAL)
            if test:
                return

    def _send_ping(self):
        """
        Sends a heartbeat to the Runpod server.
        """
        job_ids = self._mirror.get() if self._mirror is not None else None
        ping_params = {"job_id": job_ids, "runpod_version": runpod_version}

        try:
            result = self._session.get(
                self.PING_URL, params=ping_params, timeout=self.PING_INTERVAL * 2
            )

            log.debug(
                f"Heartbeat Sent | URL: {result.url} | Status: {result.status_code}"
            )
        except requests.RequestException as err:
            log.error(f"Ping Request Error: {err}, attempting to restart ping.")