# mugi_shutdown_cleanup.py — Best-effort teardown when the app or backend exits (battery / stray processes).
from __future__ import annotations

import logging
import os
import sys
import threading
import time
from typing import Any, Callable

_log = logging.getLogger(__name__)

# Exit code for a watchdog-forced exit so post-mortems can distinguish
# "shutdown hung and was force-killed" from a clean exit or a crash.
FORCED_EXIT_CODE = 70


def arm_forced_exit_watchdog(grace_seconds: float, *, reason: str = "shutdown") -> threading.Thread:
    """Guarantee the process exits within *grace_seconds* once shutdown has begun.

    June 2026 incident: a SIGTERM ran the cleanup sequence (dispatcher stopped,
    workers cancelled) but uvicorn's graceful shutdown then waited forever on
    open SSE connections — the process survived as a zombie with no HTTP
    listener and a permanently stopped dispatcher, and the Swift host never
    auto-restarted it because the PID was still alive. A daemon timer thread
    does not depend on the (possibly wedged) main thread, so it fires even
    when ``Thread.join`` / lock acquisition is blocking everything else.
    """
    deadline = time.monotonic() + max(0.5, float(grace_seconds))

    def _watch() -> None:
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                break
            time.sleep(min(remaining, 1.0))
        print(
            f"[mugi] shutdown ({reason}) did not finish within {grace_seconds:.0f}s; "
            f"forcing process exit so the host can restart a clean backend.",
            file=sys.stderr,
            flush=True,
        )
        try:
            _log.warning("shutdown watchdog fired (reason=%s); forcing exit", reason)
        except Exception:
            pass
        os._exit(FORCED_EXIT_CODE)

    t = threading.Thread(target=_watch, daemon=True, name=f"shutdown-watchdog-{reason}")
    t.start()
    return t


def run_mugi_shutdown_cleanup(*, reason: str = "shutdown") -> dict[str, Any]:
    """Stop heavy background services. Safe to call multiple times (idempotent where possible)."""
    _log.info("shutdown cleanup starting (reason=%s)", reason)
    summary: dict[str, Any] = {"reason": reason, "steps": []}

    def _step(name: str, fn: Callable[[], None]) -> None:
        try:
            fn()
            summary["steps"].append({"name": name, "ok": True})
        except Exception as e:
            _log.debug("shutdown step %s failed", name, exc_info=True)
            summary["steps"].append({"name": name, "ok": False, "error": str(e)[:200]})

    _step(
        "cron_scheduler",
        lambda: __import__("cron_scheduler", fromlist=["stop_scheduler"]).stop_scheduler(),
    )

    def _kanban_stop() -> None:
        from utils.kanban_dispatcher import stop_dispatcher
        # The dispatcher's stop() also calls kanban_storage.mark_inflight_cancelled
        # so any worker still mid-flight has its run row recorded as cancelled
        # (not crashed) and the task drops back to ready for the next process.
        stop_dispatcher()

    _step("kanban_dispatcher", _kanban_stop)

    def _evo_stop() -> None:
        from evolution.scheduler import stop_evolution_scheduler

        stop_evolution_scheduler()

    _step("evolution_scheduler", _evo_stop)

    def _ft_cancel() -> None:
        from fine_tuner import cancel_fine_tune

        cancel_fine_tune()

    _step("fine_tune_cancel", _ft_cancel)

    def _mcp_shutdown() -> None:
        from mcp_bridge.session_manager import SESSION_MANAGER

        SESSION_MANAGER.shutdown()

    _step("mcp_session_manager", _mcp_shutdown)

    def _mesh_stop() -> None:
        from mesh import get_mesh_manager

        m = get_mesh_manager()
        if m is not None:
            m.stop()

    _step("mesh", _mesh_stop)

    _step("mlx_server", lambda: __import__("app_mlx", fromlist=["stop_mlx_server"]).stop_mlx_server())

    def _llmster_stop() -> None:
        from app_llmster import stop_llmster_server

        stop_llmster_server()

    _step("llmster_server", _llmster_stop)

    def _searxng_colima() -> None:
        from core.config import get_search_config

        cfg = get_search_config()
        if not cfg.get("use_local_docker"):
            return
        from tools.searxng_docker import searxng_stop, _run_colima

        searxng_stop()
        _run_colima(["stop"], timeout=120)

    _step("searxng_colima", _searxng_colima)

    def _bg_bootout() -> None:
        import app_background_agent as bga

        if not bga._background_agent_is_installed():
            return
        ok, _msg = bga._launchctl_bootout()
        if not ok:
            # Job may already be unloaded; still counts as best-effort.
            pass

    _step("background_launch_agent", _bg_bootout)

    failed = [s["name"] for s in summary["steps"] if not s.get("ok")]
    if failed:
        _log.warning("shutdown cleanup done (reason=%s); failed steps: %s", reason, ", ".join(failed))
    else:
        _log.info("shutdown cleanup done (reason=%s); all steps ok", reason)
    return summary
