"""
Jupyter Server Manager

Manages the lifecycle of a Jupyter Server subprocess.
Handles starting, stopping, health monitoring, and token management.

Features:
- Lifecycle lock to serialize start/stop operations
- Automatic restart on unexpected server death
- Stdout/stderr draining to prevent pipe blocking

Shutdown:
- The owner (runner/KernelManager) should call stop_server() during shutdown
- A minimal atexit handler provides last-resort cleanup if owner crashes
"""

import asyncio
import atexit
import json
import logging
import os
import secrets
import shutil
import signal
import socket
import sys
import time
import weakref
from pathlib import Path
from typing import Callable, Optional

from ..structured_log import get_logger as _get_slog

import jupyter_core.paths

# Track all manager instances for last-resort cleanup on exit
_active_managers: weakref.WeakSet["JupyterServerManager"] = weakref.WeakSet()


def _atexit_cleanup() -> None:
    """
    Last-resort cleanup: terminate any orphaned Jupyter subprocesses.

    This is a safety net for edge cases (owner crashed, forgot to call stop_server).
    Normal shutdown should use the async stop_server() path via the owner.
    """
    for manager in list(_active_managers):
        try:
            manager._terminate_subprocess()
        except Exception:
            pass


# Register minimal atexit handler as safety net
atexit.register(_atexit_cleanup)

logger = logging.getLogger("skop-runner")
_log = _get_slog("jupyter_server")

# Port configuration
JUPYTER_SERVER_BASE_PORT = 18888  # Not 8888 to avoid conflict with user's Jupyter
JUPYTER_SERVER_PORT_RETRIES = 10  # Let Jupyter handle port conflicts natively

# Timeouts
STARTUP_TIMEOUT = 30.0  # seconds to wait for server to be ready
SHUTDOWN_TIMEOUT = 10.0  # seconds to wait for graceful shutdown
HEALTH_CHECK_INTERVAL = 5.0  # seconds between health checks

# Restart configuration
MAX_RESTART_ATTEMPTS = 3  # max consecutive restart attempts
RESTART_BACKOFF_BASE = 2.0  # base seconds for exponential backoff


class PortBindingError(RuntimeError):
    """Raised when server fails to bind to the requested port."""

    pass


def is_port_available(port: int, host: str = "127.0.0.1") -> bool:
    """
    Check if a port is available for binding.

    Args:
        port: Port number to check
        host: Host address to bind to

    Returns:
        True if port is available, False otherwise
    """
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            sock.bind((host, port))
            return True
    except OSError:
        return False


def find_available_port(
    base_port: int = JUPYTER_SERVER_BASE_PORT,
    max_retries: int = JUPYTER_SERVER_PORT_RETRIES,
    host: str = "127.0.0.1",
) -> int:
    """
    Find an available port, starting from base_port.

    First tries ports in range [base_port, base_port + max_retries).
    If all are busy, falls back to an OS-assigned random port.

    Args:
        base_port: Starting port number
        max_retries: Number of sequential ports to try
        host: Host address to bind to

    Returns:
        Available port number
    """
    # Try preferred port range first
    for port in range(base_port, base_port + max_retries):
        if is_port_available(port, host):
            return port

    # All preferred ports busy, fall back to OS-assigned random port
    logger.info(
        f"Ports {base_port}-{base_port + max_retries - 1} all busy, "
        "falling back to random OS-assigned port"
    )

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind((host, 0))  # Port 0 = let OS assign
        return sock.getsockname()[1]


class JupyterServerManager:
    """
    Manages Jupyter Server subprocess lifecycle.

    Responsibilities:
    - Start Jupyter Server on an available port
    - Generate and manage authentication token
    - Monitor server health with automatic restart
    - Graceful shutdown with cleanup
    - Reads Jupyter's runtime file for actual port/token (secure permissions)
    - Drain stdout/stderr to prevent pipe blocking

    Thread Safety:
    - All lifecycle operations (start/stop/restart) are serialized via _lifecycle_lock
    - Safe to call ensure_running() concurrently from multiple coroutines
    """

    # Persistence format version — bump when schema changes
    _PERSISTENCE_VERSION = 1

    def __init__(
        self,
        on_server_died: Optional[Callable[[], None]] = None,
        on_server_restarted: Optional[Callable[[], None]] = None,
        persistence_dir: Optional[str] = None,
    ):
        """
        Initialize the Jupyter Server Manager.

        Args:
            on_server_died: Optional callback invoked when server dies unexpectedly
                           (after restart attempts exhausted). Can be used to notify frontend.
            on_server_restarted: Optional callback invoked after successful automatic restart.
            persistence_dir: Optional directory for persisting server info to disk.
                If provided, PID/port/token are saved to ``{persistence_dir}/server.json``
                and checked on startup to reuse a surviving server.
        """
        self.port: Optional[int] = None
        self.token: Optional[str] = None
        self.base_url: Optional[str] = None
        self._process: Optional[asyncio.subprocess.Process] = None
        self._monitor_task: Optional[asyncio.Task] = None
        self._stdout_task: Optional[asyncio.Task] = None
        self._stderr_task: Optional[asyncio.Task] = None
        self._runtime_file: Optional[Path] = None  # Jupyter's runtime file (read-only)
        self._root_dir: Optional[str] = None
        self._stderr_buffer: list[str] = []  # Capture stderr for error diagnosis

        # Lifecycle control - lazily initialized to bind to the correct event loop.
        # Manager may be instantiated before asyncio.run(), so we defer creation
        # until first use within the running loop.
        self.__lifecycle_lock: Optional[asyncio.Lock] = None
        self.__ready_event: Optional[asyncio.Event] = None
        self._restart_count = 0
        self._shutting_down = False
        self._on_server_died = on_server_died
        self._on_server_restarted = on_server_restarted

        # PID of a reused server (when we don't own the subprocess)
        self._reused_pid: Optional[int] = None

        # Persistence
        self._persistence_path: Optional[Path] = None
        if persistence_dir:
            self._persistence_path = Path(persistence_dir) / "server.json"

        # Register for cleanup on runner exit
        _active_managers.add(self)

    @property
    def _lifecycle_lock(self) -> asyncio.Lock:
        """Lazily create lock bound to the current running event loop."""
        if self.__lifecycle_lock is None:
            self.__lifecycle_lock = asyncio.Lock()
        return self.__lifecycle_lock

    @property
    def _ready_event(self) -> asyncio.Event:
        """Lazily create event bound to the current running event loop."""
        if self.__ready_event is None:
            self.__ready_event = asyncio.Event()
        return self.__ready_event

    def is_running(self) -> bool:
        """Check if server process is alive."""
        if self._process is not None and self._process.returncode is None:
            return True
        # Reused server from previous runner session — check PID is still alive
        if self._reused_pid is not None:
            return self._is_pid_alive(self._reused_pid)
        return False

    # ------------------------------------------------------------------
    # Disk persistence
    # ------------------------------------------------------------------

    def _save_server_info(self) -> None:
        """Save server PID/port/token to disk for runner restart survival."""
        if self._persistence_path is None or not self._process:
            return

        data = {
            "version": self._PERSISTENCE_VERSION,
            "pid": self._process.pid,
            "port": self.port,
            "token": self.token,
            "base_url": self.base_url,
            "root_dir": self._root_dir,
        }

        try:
            self._persistence_path.parent.mkdir(parents=True, exist_ok=True)
            tmp_path = self._persistence_path.with_suffix(".json.tmp")
            tmp_path.write_text(json.dumps(data, indent=2))
            os.replace(str(tmp_path), str(self._persistence_path))
            from skop_runner.compat import secure_file_permissions
            secure_file_permissions(str(self._persistence_path))
            logger.debug(f"Saved server info to {self._persistence_path}")
        except Exception as e:
            logger.warning(f"Failed to persist server info: {e}")

    def _load_server_info(self) -> Optional[dict]:
        """
        Load saved server info from disk.

        Returns:
            Dict with pid/port/token/base_url/root_dir if file exists and is valid,
            None otherwise.
        """
        if self._persistence_path is None or not self._persistence_path.exists():
            return None

        try:
            raw = json.loads(self._persistence_path.read_text())
        except (json.JSONDecodeError, IOError) as e:
            logger.warning(f"Failed to read server info file: {e}")
            return None

        if not isinstance(raw, dict) or raw.get("version") != self._PERSISTENCE_VERSION:
            logger.warning(f"Server info file has wrong version, ignoring")
            return None

        # Validate required fields
        required = ["pid", "port", "token", "base_url", "root_dir"]
        if not all(raw.get(k) is not None for k in required):
            logger.warning("Server info file missing required fields")
            return None

        return raw

    def _delete_server_info(self) -> None:
        """Delete the server info file from disk."""
        if self._persistence_path and self._persistence_path.exists():
            try:
                self._persistence_path.unlink()
                logger.debug(f"Deleted server info file: {self._persistence_path}")
            except Exception as e:
                logger.warning(f"Failed to delete server info file: {e}")

    def _is_pid_alive(self, pid: int) -> bool:
        """Check if a process with the given PID is still running."""
        from skop_runner.compat import is_pid_alive
        return is_pid_alive(pid)

    async def _try_reuse_server(self, root_dir: Optional[str]) -> Optional[str]:
        """
        Try to reuse a surviving Jupyter Server from a previous runner session.

        Checks if the saved PID is alive, the API responds with our token,
        and the root_dir matches.

        Returns:
            Authentication token if server was reused, None otherwise.
        """
        saved = self._load_server_info()
        if saved is None:
            return None

        pid = saved["pid"]
        saved_root = saved["root_dir"]
        effective_root = root_dir or self._root_dir or os.getcwd()

        # Root dir must match
        if saved_root != effective_root:
            logger.info(
                f"Saved server has different root ({saved_root}), "
                f"not reusing (need {effective_root})"
            )
            return None

        # PID must be alive
        if not self._is_pid_alive(pid):
            logger.info(f"Saved server PID {pid} is no longer running")
            self._delete_server_info()
            return None

        # Verify it's actually our Jupyter Server by hitting the API
        import httpx

        base_url = saved["base_url"]
        token = saved["token"]
        try:
            async with httpx.AsyncClient() as client:
                resp = await client.get(
                    f"{base_url}/api/status",
                    headers={"Authorization": f"token {token}"},
                    timeout=5.0,
                )
                if resp.status_code != 200:
                    logger.info(
                        f"Saved server at {base_url} returned status {resp.status_code}, not reusing"
                    )
                    self._delete_server_info()
                    return None
        except Exception as e:
            logger.info(f"Failed to reach saved server at {base_url}: {e}")
            self._delete_server_info()
            return None

        # Server is alive and healthy — restore state
        self.port = saved["port"]
        self.token = token
        self.base_url = base_url
        self._root_dir = saved_root
        self._reused_pid = pid
        self._shutting_down = False
        self._restart_count = 0

        # Start health monitor for the reused server
        if self._monitor_task and not self._monitor_task.done():
            self._monitor_task.cancel()
            try:
                await self._monitor_task
            except asyncio.CancelledError:
                pass
        self._monitor_task = asyncio.create_task(self._monitor_health())

        logger.info(
            f"Reusing existing Jupyter Server (PID {pid}) on {base_url}"
        )
        return token

    async def detach(self) -> None:
        """
        Detach from the server without stopping it.

        Cancels health monitoring, closes drain tasks, and removes from
        atexit cleanup.  The server process continues running for reuse
        by the next runner session.
        """
        # Cancel health monitor
        if self._monitor_task and not self._monitor_task.done():
            self._monitor_task.cancel()
            try:
                await self._monitor_task
            except asyncio.CancelledError:
                pass
            self._monitor_task = None

        # Cancel output drain tasks (pipes close when process exits)
        await self._cancel_drain_tasks()

        # Remove from atexit cleanup so _terminate_subprocess is not called
        _active_managers.discard(self)

    def _terminate_subprocess(self) -> None:
        """
        Send SIGTERM to subprocess. Last-resort safety net for atexit.

        This is intentionally minimal - just sends terminate signal and returns.
        The proper shutdown path is stop_server() which the owner should call.

        Note: Does NOT terminate reused servers (they were not started by us).
        """
        if not self._process:
            return
        if self._process.returncode is not None:
            return
        try:
            self._process.terminate()
        except (OSError, ProcessLookupError):
            pass  # Process already dead

    async def start_server(self, root_dir: Optional[str] = None) -> str:
        """
        Start Jupyter Server subprocess.

        This method is serialized via lifecycle lock to prevent race conditions.

        Args:
            root_dir: Working directory for Jupyter Server (defaults to cwd)

        Returns:
            Generated authentication token

        Raises:
            RuntimeError: If server fails to start or is already running
        """
        async with self._lifecycle_lock:
            return await self._start_server_unlocked(root_dir)

    async def _start_server_unlocked(
        self,
        root_dir: Optional[str] = None,
        *,
        _skip_monitor: bool = False,
    ) -> str:
        """
        Internal start implementation (caller must hold _lifecycle_lock).

        Args:
            root_dir: Working directory for Jupyter Server
            _skip_monitor: If True, don't create a new monitor task (used when
                          called from within the monitor's restart logic).
        """
        _t0 = time.perf_counter()
        if self.is_running():
            logger.warning("Jupyter Server already running, returning existing token")
            return self.token

        # Clear ready event - will be set after successful startup
        self._ready_event.clear()
        self._shutting_down = False

        # Use local variable for root_dir during startup attempt.
        # Only commit to self._root_dir after startup succeeds.
        # This ensures failed startups don't leave stale _root_dir that would
        # cause subsequent workspace.setRoot calls with the same path to skip restart.
        effective_root_dir = root_dir or self._root_dir or os.getcwd()

        # Generate fresh token for security - prevents user config from disabling auth
        self.token = secrets.token_hex(32)
        logger.debug("Generated new authentication token")

        # Find an available port
        requested_port = find_available_port()
        logger.info(f"Found available port {requested_port} for Jupyter Server")

        # Resolve the full path to `jupyter` so it works even when the
        # virtualenv isn't activated (e.g. running via /path/to/venv/bin/skop-runner).
        jupyter_bin = shutil.which("jupyter")
        if jupyter_bin is None:
            # Fall back to looking next to the running Python executable.
            # Use shutil.which with explicit path so PATHEXT is respected
            # on Windows (finds jupyter.exe, jupyter.cmd, etc.)
            bin_dir = str(Path(sys.executable).parent)
            jupyter_bin = shutil.which("jupyter", path=bin_dir)
            if jupyter_bin is None:
                raise RuntimeError(
                    "jupyter command not found. Is jupyter-server installed? "
                    "Run: pip install jupyter-server"
                )

        cmd = [
            jupyter_bin,
            "server",
            "--ServerApp.ip=127.0.0.1",
            f"--ServerApp.port={requested_port}",
            "--ServerApp.port_retries=3",
            "--ServerApp.allow_origin=*",
            "--ServerApp.disable_check_xsrf=True",
            f"--ServerApp.root_dir={effective_root_dir}",
            "--ServerApp.log_level=WARN",
            "--KernelRestarter.restart_limit=0",
            "--MappingKernelManager.cull_idle_timeout=600",
            "--MappingKernelManager.cull_interval=60",
            "--MappingKernelManager.cull_connected=False",
            "--no-browser",
        ]

        # Pass token via JUPYTER_TOKEN environment variable for security.
        # This keeps the token off the process command line (visible via ps/proc).
        # Note: User config files have higher precedence than env vars in traitlets,
        # so users with c.ServerApp.token="" in their config will see a clear error
        # from _wait_for_ready() explaining how to fix the conflict.
        env = os.environ.copy()
        env["JUPYTER_TOKEN"] = self.token

        # Add our custom kernelspec directory to Jupyter's search path
        # This allows ephemeral kernelspecs registered via register_ephemeral_kernelspec()
        # JUPYTER_PATH accepts multiple paths separated by os.pathsep
        skop_data_dir = Path.home() / ".skop"
        existing_path = env.get("JUPYTER_PATH", "")
        if existing_path:
            env["JUPYTER_PATH"] = f"{skop_data_dir}{os.pathsep}{existing_path}"
        else:
            env["JUPYTER_PATH"] = str(skop_data_dir)

        try:
            self._process = await asyncio.create_subprocess_exec(
                *cmd,
                env=env,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
        except FileNotFoundError:
            raise RuntimeError(
                f"jupyter command not found at {jupyter_bin}. "
                "Is jupyter-server installed? Run: pip install jupyter-server"
            )

        # Clear stderr buffer for this startup attempt
        self._stderr_buffer.clear()

        # Start output draining tasks immediately to prevent pipe blocking
        self._stdout_task = asyncio.create_task(
            self._drain_output(self._process.stdout, "stdout")
        )
        self._stderr_task = asyncio.create_task(
            self._drain_output(self._process.stderr, "stderr")
        )

        try:
            # Wait for server to be ready (reads Jupyter's runtime file)
            if not await self._wait_for_ready():
                await self._cleanup_failed_start()
                raise RuntimeError(
                    f"Jupyter Server failed to start within {STARTUP_TIMEOUT}s. "
                    "Check that jupyter-server is installed correctly."
                )
        except PortBindingError:
            await self._cleanup_failed_start()
            raise
        except asyncio.CancelledError:
            logger.warning("Server startup cancelled, cleaning up...")
            await self._cleanup_failed_start()
            raise

        # Startup succeeded - now commit the root_dir
        # This ensures _root_dir only reflects a successfully started server
        self._root_dir = effective_root_dir
        self._reused_pid = None  # We now own a subprocess; clear any stale reused PID

        # Bug fix #2: Don't create monitor if called from monitor's restart logic
        # Bug fix #4: Cancel existing monitor to prevent duplicate monitors
        if not _skip_monitor:
            if self._monitor_task and not self._monitor_task.done():
                self._monitor_task.cancel()
                try:
                    await self._monitor_task
                except asyncio.CancelledError:
                    pass
            self._monitor_task = asyncio.create_task(self._monitor_health())

        # Reset restart count on successful start
        self._restart_count = 0

        # Signal that server is fully ready with valid credentials
        self._ready_event.set()

        self._save_server_info()
        _log.info(
            "jupyter_server.start",
            port=self.port,
            duration_ms=round((time.perf_counter() - _t0) * 1000),
        )
        return self.token

    async def _drain_output(
        self, stream: asyncio.StreamReader, stream_name: str
    ) -> None:
        """
        Continuously drain output from subprocess stream into logger.

        This prevents the subprocess from blocking on full pipes and ensures
        diagnostic information is captured in logs. Stderr is also buffered
        for error diagnosis (e.g., detecting port binding failures).

        Args:
            stream: The stdout or stderr stream to drain
            stream_name: 'stdout' or 'stderr' for log prefixing
        """
        try:
            while True:
                line = await stream.readline()
                if not line:
                    break
                decoded = line.decode("utf-8", errors="replace").rstrip()
                if decoded:
                    if stream_name == "stderr":
                        logger.warning(f"[jupyter-server] {decoded}")
                        # Buffer stderr for error diagnosis (keep last 50 lines)
                        self._stderr_buffer.append(decoded)
                        if len(self._stderr_buffer) > 50:
                            self._stderr_buffer.pop(0)
                    else:
                        logger.debug(f"[jupyter-server] {decoded}")
        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.debug(f"Output drain for {stream_name} ended: {e}")

    async def _wait_for_ready(self) -> bool:
        """
        Wait for server to be ready by checking Jupyter's runtime file.

        Jupyter Server writes its runtime info (including actual port and token)
        to ~/.local/share/jupyter/runtime/jpserver-<pid>.json (or platform equivalent).
        This file has secure permissions set by Jupyter.

        Readiness is determined purely by the runtime file existing with a valid,
        matching token. We skip HTTP probing during startup because:
        - If Jupyter wrote the file with our token, it's running
        - HTTP probes can fail due to proxy settings, firewall rules, timing
        - The health monitor handles HTTP-level verification after startup
        """
        start_time = asyncio.get_running_loop().time()
        runtime_dir = Path(jupyter_core.paths.jupyter_runtime_dir())
        expected_file = runtime_dir / f"jpserver-{self._process.pid}.json"

        while asyncio.get_running_loop().time() - start_time < STARTUP_TIMEOUT:
            # Check if process died
            if self._process.returncode is not None:
                logger.error("Jupyter Server died during startup")
                # Check stderr for port binding errors
                stderr_text = "\n".join(self._stderr_buffer).lower()
                if (
                    "address already in use" in stderr_text
                    or "eaddrinuse" in stderr_text
                    or "port is already allocated" in stderr_text
                    or "bind" in stderr_text and "failed" in stderr_text
                ):
                    raise PortBindingError(
                        f"Port {self.port} is already in use. "
                        "Another process may be using this port."
                    )
                return False

            # Check for Jupyter's runtime file
            if expected_file.exists():
                try:
                    data = json.loads(expected_file.read_text())

                    # Verify runtime file has token and it matches what we set.
                    # This ensures our token was applied and not overridden by user config.
                    runtime_token = data.get("token")
                    if not runtime_token:
                        logger.debug("Runtime file exists but token not yet written")
                        await asyncio.sleep(0.5)
                        continue

                    if runtime_token != self.token:
                        # User's Jupyter config is overriding our token
                        logger.error(
                            "Jupyter config is overriding the token. "
                            "Remove token settings from ~/.jupyter/jupyter_server_config.py"
                        )
                        return False

                    actual_port = data.get("port")
                    hostname = data.get("hostname") or "127.0.0.1"

                    # ServerApp.base_url controls the prefix for every API endpoint.
                    # Use it directly instead of parsing the 'url' field which may
                    # include UI paths (/lab) or query strings (?token=...).
                    base_path = data.get("base_url", "/") or "/"
                    if not base_path.startswith("/"):
                        base_path = f"/{base_path}"
                    base_path = base_path.rstrip("/")
                    if base_path == "/":
                        base_path = ""

                    secure_flag = data.get("secure")
                    if secure_flag is None:
                        url_value = data.get("url")
                        secure_flag = isinstance(url_value, str) and url_value.startswith(
                            "https://"
                        )

                    scheme = "https" if secure_flag else "http"

                    self.base_url = f"{scheme}://{hostname}:{actual_port}{base_path}"
                    self.port = actual_port
                    logger.info(f"Jupyter Server bound to port {actual_port}")

                    self._runtime_file = expected_file
                    logger.debug(f"Found Jupyter runtime file: {expected_file}")

                    # Runtime file valid with matching token - server is ready
                    # Health monitor will handle HTTP-level verification
                    return True

                except (json.JSONDecodeError, IOError, KeyError) as e:
                    logger.debug(f"Runtime file not ready yet: {e}")

            await asyncio.sleep(0.5)

        return False

    async def _cleanup_failed_start(self):
        """Clean up after failed startup."""
        await self._cancel_drain_tasks()

        if self._process:
            if self._process.returncode is None:
                self._process.terminate()
                try:
                    await asyncio.wait_for(self._process.wait(), timeout=5.0)
                except asyncio.TimeoutError:
                    if self._process.returncode is None:
                        self._process.kill()
                    await self._process.wait()

        self._process = None
        self.port = None
        self.token = None
        self.base_url = None

    async def _cancel_drain_tasks(self):
        """Cancel stdout/stderr drain tasks."""
        for task in [self._stdout_task, self._stderr_task]:
            if task and not task.done():
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
        self._stdout_task = None
        self._stderr_task = None

    async def _monitor_health(self):
        """
        Background task to monitor server health and handle automatic restart.

        When the server dies unexpectedly:
        1. Attempts automatic restart with exponential backoff
        2. After MAX_RESTART_ATTEMPTS, gives up and notifies via callback
        """
        while True:
            await asyncio.sleep(HEALTH_CHECK_INTERVAL)

            if self._shutting_down:
                break

            if not self.is_running():
                _log.error(
                    "jupyter_server.health_check_failed",
                    port=self.port,
                    error="process_died",
                )

                if self._restart_count < MAX_RESTART_ATTEMPTS:
                    self._restart_count += 1
                    backoff = RESTART_BACKOFF_BASE ** self._restart_count
                    _log.info(
                        "jupyter_server.restart_attempt",
                        attempt_number=self._restart_count,
                        max_attempts=MAX_RESTART_ATTEMPTS,
                    )
                    await asyncio.sleep(backoff)

                    try:
                        async with self._lifecycle_lock:
                            await self._cancel_drain_tasks()
                            self._process = None
                            self._reused_pid = None
                            await self._start_server_unlocked(
                                self._root_dir,
                                _skip_monitor=True,
                            )
                            logger.info("Jupyter Server restarted successfully")
                            if self._on_server_restarted:
                                try:
                                    self._on_server_restarted()
                                except Exception as e:
                                    logger.error(f"on_server_restarted callback failed: {e}")
                            continue
                    except Exception as e:
                        logger.error(f"Restart attempt {self._restart_count} failed: {e}")
                else:
                    _log.error(
                        "jupyter_server.restart_exhausted",
                        attempts=MAX_RESTART_ATTEMPTS,
                    )
                    if self._on_server_died:
                        try:
                            self._on_server_died()
                        except Exception as e:
                            logger.error(f"on_server_died callback failed: {e}")
                    break

    async def stop_server(self) -> None:
        """
        Gracefully stop the server.

        This method is serialized via lifecycle lock to prevent race conditions.
        """
        async with self._lifecycle_lock:
            await self._stop_server_unlocked()

    async def _stop_server_unlocked(self) -> None:
        """Internal stop implementation (caller must hold _lifecycle_lock)."""
        _t0 = time.perf_counter()
        _port = self.port
        # Signal that we're shutting down intentionally (not a crash)
        self._shutting_down = True
        # Clear ready event - server is no longer ready
        self._ready_event.clear()

        # Cancel monitor task first
        if self._monitor_task:
            self._monitor_task.cancel()
            try:
                await self._monitor_task
            except asyncio.CancelledError:
                pass
            self._monitor_task = None

        # Cancel output drain tasks
        await self._cancel_drain_tasks()

        # Handle reused server (no subprocess reference — kill by PID)
        if self._reused_pid is not None and self._process is None:
            pid = self._reused_pid
            logger.info(f"Stopping reused Jupyter Server (PID {pid})...")
            from skop_runner.compat import terminate_process
            await asyncio.to_thread(terminate_process, pid, SHUTDOWN_TIMEOUT)
            self._reused_pid = None
            self.port = None
            self.token = None
            self.base_url = None
            self._root_dir = None
            self._runtime_file = None
            self._restart_count = 0
            self._delete_server_info()
            logger.info("Reused Jupyter Server stopped")
            return

        if not self._process:
            # Process already gone, just reset state
            self.port = None
            self.token = None
            self.base_url = None
            self._root_dir = None
            self._runtime_file = None
            self._restart_count = 0
            self._reused_pid = None
            return

        logger.info("Stopping Jupyter Server...")

        # Only terminate if process is still alive (returncode is None)
        # After a crash, returncode is set and terminate() would raise ProcessLookupError
        if self._process.returncode is None:
            # Try graceful shutdown first
            self._process.terminate()
            try:
                await asyncio.wait_for(self._process.wait(), timeout=SHUTDOWN_TIMEOUT)
                logger.info("Jupyter Server stopped gracefully")
            except asyncio.TimeoutError:
                logger.warning("Jupyter Server didn't stop gracefully, killing")
                if self._process.returncode is None:  # Check again before kill
                    self._process.kill()
                await self._process.wait()
        else:
            logger.info("Jupyter Server process already exited")

        # Cleanup process state
        self._process = None
        self._reused_pid = None
        self.port = None
        self.token = None
        self.base_url = None
        self._root_dir = None
        self._runtime_file = None
        self._restart_count = 0
        self._delete_server_info()

        _log.info(
            "jupyter_server.stop",
            port=_port,
            duration_ms=round((time.perf_counter() - _t0) * 1000),
        )

    async def get_server_info(self) -> dict:
        """Return server URL, token, status."""
        pid = None
        if self._process:
            pid = self._process.pid
        elif self._reused_pid is not None:
            pid = self._reused_pid
        return {
            "url": self.base_url,
            "port": self.port,
            "token": self.token,
            "running": self.is_running(),
            "root_dir": self._root_dir,
            "pid": pid,
        }

    async def ensure_running(self, root_dir: Optional[str] = None) -> str:
        """
        Ensure server is running, starting it if needed.

        This is the main entry point for lazy initialization.
        Safe to call concurrently from multiple coroutines.
        Handles cancellation safely (cleans up subprocess if cancelled mid-start).

        Args:
            root_dir: Working directory for Jupyter Server

        Returns:
            Authentication token
        """
        # Quick check without lock for common case (already running and ready)
        if self.is_running() and self._ready_event.is_set():
            return self.token

        # Need to acquire lock to potentially start
        async with self._lifecycle_lock:
            # Re-check after acquiring lock (another coroutine may have started it)
            if self.is_running() and self._ready_event.is_set():
                return self.token

            # Try to reuse a surviving server from a previous runner session
            token = await self._try_reuse_server(root_dir)
            if token is not None:
                self._ready_event.set()
                return token

            # _start_server_unlocked handles its own cancellation cleanup internally
            return await self._start_server_unlocked(root_dir)

    async def restart_server(self) -> str:
        """
        Restart the Jupyter Server.

        Stops the server if running, then starts it again.
        Useful for recovering from errors or applying configuration changes.

        Returns:
            New authentication token
        """
        async with self._lifecycle_lock:
            # Capture root_dir before stopping (stop clears it)
            saved_root_dir = self._root_dir
            if self.is_running():
                await self._stop_server_unlocked()
            return await self._start_server_unlocked(saved_root_dir)

    async def set_root(self, root_dir: str) -> str:
        """
        Change the server's root directory atomically.

        If the server is running with a different root, it will be stopped and
        restarted with the new root. If the restart fails, the server remains
        stopped (no partial state).

        This method is the safe way to change project roots - it ensures:
        - Server is only restarted if root actually changed
        - _root_dir only updates after successful startup
        - Failed startups can be retried with the same path

        Args:
            root_dir: New root directory (will be resolved to absolute path)

        Returns:
            Authentication token

        Raises:
            RuntimeError: If server fails to start with new root
        """
        from pathlib import Path

        resolved_root = str(Path(root_dir).expanduser().resolve())

        async with self._lifecycle_lock:
            # Check if we need to do anything
            if self._root_dir == resolved_root and self.is_running():
                logger.debug(f"Server already running with root {resolved_root}")
                return self.token

            # Stop if running (regardless of current root)
            if self.is_running():
                logger.info(f"Stopping server to change root to {resolved_root}")
                await self._stop_server_unlocked()

            # Start with new root - _root_dir only updated on success
            return await self._start_server_unlocked(resolved_root)

    async def register_ephemeral_kernelspec(
        self,
        python_path: str,
        env_name: Optional[str] = None,
    ) -> str:
        """
        Register an ephemeral kernelspec for a custom Python executable.

        Creates a temporary kernelspec in ~/.skop/kernels/ that points to
        the specified Python executable. This allows using custom venvs/condas
        that don't have kernelspecs installed.

        Args:
            python_path: Absolute path to Python executable
            env_name: Optional display name for the environment

        Returns:
            The kernelspec name to use when starting a kernel

        Note:
            The kernelspec is persistent in ~/.skop/kernels/ but prefixed
            with 'skop-' to distinguish from system kernelspecs.
        """
        import hashlib

        # Generate a unique but stable name based on the Python path
        path_hash = hashlib.sha256(python_path.encode()).hexdigest()[:12]
        spec_name = f"skop-{path_hash}"

        # Determine display name
        if env_name:
            display_name = f"Python ({env_name})"
        else:
            display_name = f"Python ({Path(python_path).parent.parent.name})"

        # Create kernelspec directory
        # Jupyter expects kernels at {JUPYTER_PATH}/kernels/{name}/kernel.json
        kernelspec_dir = Path.home() / ".skop" / "kernels" / spec_name
        kernelspec_dir.mkdir(parents=True, exist_ok=True)

        # Write kernel.json
        kernel_json = {
            "argv": [
                python_path,
                "-m",
                "ipykernel_launcher",
                "-f",
                "{connection_file}",
            ],
            "display_name": display_name,
            "language": "python",
            "metadata": {
                "skop_managed": True,
                "python_path": python_path,
            },
        }

        kernel_json_path = kernelspec_dir / "kernel.json"
        kernel_json_path.write_text(json.dumps(kernel_json, indent=2))

        logger.info(f"Registered ephemeral kernelspec: {spec_name} -> {python_path}")

        return spec_name
