"""Register atexit and signal handlers so workload completion is always signaled.

When 'import perpetual_compute' runs, this module is called once. It registers:
  - An atexit hook: if no explicit done()/fail() was called, infers status from
    sys.exc_info() (exception in flight = "failed", clean = "completed").
  - SIGTERM handler: marks as "failed" (spot interruption signal), writes status,
    and re-raises the original SIGTERM so the OS continues normal shutdown.

Migration safety: if CRIU_MIGRATION_IN_PROGRESS=1 is set in the environment,
all hooks are skipped. The migration engine sets this variable before checkpointing
so a CRIU freeze does not look like workload completion.
"""

from __future__ import annotations

import atexit
import os
import signal
import sys
from typing import Any

from perpetual_compute._status import is_written, write_status

_hooks_registered: bool = False


def _atexit_hook() -> None:
    if os.environ.get("CRIU_MIGRATION_IN_PROGRESS") == "1":
        return
    if is_written():
        return
    # Determine status from whether an exception is currently active.
    exc_type, _, _ = sys.exc_info()
    if exc_type is not None:
        write_status("failed", reason=f"Unhandled exception: {exc_type.__name__}")
    else:
        write_status("completed")


def _sigterm_handler(signum: int, frame: Any) -> None:
    if not os.environ.get("CRIU_MIGRATION_IN_PROGRESS") == "1":
        if not is_written():
            write_status("failed", reason="SIGTERM received")
    # Re-raise as default SIGTERM so the process actually terminates.
    signal.signal(signal.SIGTERM, signal.SIG_DFL)
    os.kill(os.getpid(), signal.SIGTERM)


def register() -> None:
    global _hooks_registered
    if _hooks_registered:
        return
    _hooks_registered = True
    atexit.register(_atexit_hook)
    try:
        signal.signal(signal.SIGTERM, _sigterm_handler)
    except (OSError, ValueError):
        # Cannot set signal handlers in non-main threads — safe to ignore.
        pass
