from logging import Logger

import atexit
import signal
import threading
from types import FrameType
from typing import Callable, Optional, Union, cast

Handler = Union[int, Callable[[int, Optional[FrameType]], None], None]


def install_frisk_shutdown(
    shutdown_fn: Callable[[float], None],
    *,
    timeout_s: float,
    install_signals: bool,
    logger: Logger,
) -> None:
    """Install idempotent atexit/signal shutdown hooks for Frisk instances."""
    state = {"has_run": False}  # closure state, easy to reset in tests

    def run(reason: str) -> None:
        if state["has_run"]:
            return
        state["has_run"] = True

        # Run shutdown best-effort with a wall-clock budget
        done = threading.Event()

        def worker() -> None:
            try:
                shutdown_fn(timeout_s)
            except Exception as e:
                logger.error(f"Error running shutdown handler: {e}")
                pass
            finally:
                done.set()

        t = threading.Thread(target=worker, name="frisk-shutdown", daemon=True)
        t.start()
        done.wait(timeout=timeout_s)

    atexit.register(lambda: run("atexit"))

    if not install_signals:
        return

    def chain(signum: int, prev: Handler):
        def handler(sig: int, frame: Optional[FrameType]) -> None:
            run(f"signal:{sig}")
            # chain previous
            if prev is signal.SIG_IGN:
                return
            if prev is signal.SIG_DFL:
                if sig == signal.SIGINT:
                    raise KeyboardInterrupt
                return
            if prev is None or isinstance(prev, int):
                return
            cast(Callable[[int, Optional[FrameType]], None], prev)(sig, frame)

        return handler

    for s in (signal.SIGINT, signal.SIGTERM):
        prev = signal.getsignal(s)
        signal.signal(s, chain(s, prev))
