"""
runpod | serverless | worker_loop.py
Called to convert a container into a worker pod for the runpod serverless platform.
"""

import asyncio
import os
import threading
from http.server import HTTPServer, BaseHTTPRequestHandler
from typing import Any, Dict

from gmi_ieops.runpod.serverless.modules import rp_logger, rp_local, rp_ping, rp_scale

log = rp_logger.RunPodLogger()
heartbeat = rp_ping.Heartbeat()


# ---------------------------------------------------------------------------
# Health Server — lightweight HTTP endpoint for K8s liveness/readiness probes
# ---------------------------------------------------------------------------

class _HealthHandler(BaseHTTPRequestHandler):
    """Minimal HTTP handler that responds to GET /health with 200 OK."""

    def do_GET(self):
        if self.path == "/health":
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(b'{"status":"ok"}')
        else:
            self.send_response(404)
            self.end_headers()

    def log_message(self, format, *args):
        pass  # suppress access logs to avoid noise


def _start_health_server(port: int = 8080) -> None:
    """Start a background daemon thread serving GET /health."""
    try:
        server = HTTPServer(("0.0.0.0", port), _HealthHandler)
        t = threading.Thread(target=server.serve_forever, daemon=True)
        t.start()
        log.info(f"Health server started on port {port}")
    except OSError as e:
        log.warn(f"Could not start health server on port {port}: {e}")


def _is_local(config) -> bool:
    """Returns True if the worker is running locally, False otherwise."""
    if config["rp_args"].get("test_input", None):
        return True

    if os.environ.get("RUNPOD_WEBHOOK_GET_JOB", None) is None:
        return True

    return False


# ------------------------- Main Worker Running Loop ------------------------- #
def run_worker(config: Dict[str, Any]) -> None:
    """
    Starts the worker loop for multi-processing.

    This function is called when the worker is running on Runpod. This function
    starts a loop that runs indefinitely until the worker is killed.

    Args:
        config (Dict[str, Any]): Configuration parameters for the worker.
    """
    # Start pinging Runpod to show that the worker is alive.
    heartbeat.start_ping()

    # Start health server for K8s probes (ieops-v2 addition).
    health_port = int(os.environ.get("HEALTH_PORT", "8080"))
    _start_health_server(port=health_port)

    # Create a JobScaler responsible for adjusting the concurrency
    job_scaler = rp_scale.JobScaler(config)
    job_scaler.start()


def main(config: Dict[str, Any]) -> None:
    """
    Checks if the worker is running locally or on Runpod.
    If running locally, the test job is run and the worker exits.
    If running on Runpod, the worker loop is created.
    """
    if _is_local(config):
        asyncio.run(rp_local.run_local(config))

    else:
        run_worker(config)
