"""
Session Manager for notebook-to-session-to-kernel mapping.

Provides session tracking using the KernelBackend protocol:
- Each notebook gets its own session and kernel (default behavior)
- Optional kernel sharing between notebooks (via connect_to_kernel)
- Session persistence across requests within a runner session
"""

import asyncio
import json
import logging
import os
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple

if TYPE_CHECKING:
    from .kernel_backend import KernelBackend, KernelConnectionInfo

from .structured_log import get_logger as _get_slog

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


@dataclass
class NotebookSession:
    """Represents a session for a notebook with its associated kernel."""

    session_id: str
    kernel_id: str
    notebook_path: str
    created_at: datetime
    kernel_name: str = "python3"
    python_path: Optional[str] = None
    env_name: Optional[str] = None


class SessionManager:
    """
    Manages notebook sessions and their kernel associations.

    Provides the mapping between notebook paths and kernel sessions,
    enabling automatic kernel lifecycle management tied to notebooks.

    Thread-safe for concurrent session operations via asyncio.Lock.
    """

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

    def __init__(
        self,
        kernel_backend: "KernelBackend",
        persistence_dir: Optional[str] = None,
    ):
        """
        Initialize the session manager.

        Args:
            kernel_backend: KernelBackend for starting/stopping kernels.
            persistence_dir: Optional directory for persisting sessions to disk.
                If provided, sessions are saved to ``{persistence_dir}/sessions.json``
                after every mutation and can be restored via ``load_from_disk()``.
        """
        self._kernel_backend = kernel_backend
        self._sessions: Dict[str, NotebookSession] = {}  # notebook_path -> session
        self._lock: Optional[asyncio.Lock] = None
        self._persistence_path: Optional[Path] = None
        if persistence_dir:
            self._persistence_path = Path(persistence_dir) / "sessions.json"

    @property
    def _async_lock(self) -> asyncio.Lock:
        """Get or create the lock (lazy to avoid event loop binding in __init__)."""
        if self._lock is None:
            self._lock = asyncio.Lock()
        return self._lock

    def _normalize_path(self, notebook_path: str) -> str:
        """
        Normalize notebook path for consistent lookup.

        Converts to relative path from project root if possible,
        otherwise uses the path as-is.
        """
        # Remove leading/trailing whitespace
        path = notebook_path.strip()

        # Normalize path separators and resolve relative components
        path = str(Path(path))

        return path

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

    def _save_to_disk(self) -> None:
        """
        Persist current sessions to disk as JSON.

        Called after every mutation.  Writes atomically (write-to-temp then
        rename) so a crash mid-write won't corrupt the file.
        """
        if self._persistence_path is None:
            return

        data = {
            "version": self._PERSISTENCE_VERSION,
            "sessions": {
                path: {
                    "session_id": s.session_id,
                    "kernel_id": s.kernel_id,
                    "notebook_path": s.notebook_path,
                    "created_at": s.created_at.isoformat(),
                    "kernel_name": s.kernel_name,
                    "python_path": s.python_path,
                    "env_name": s.env_name,
                }
                for path, s in self._sessions.items()
            },
        }

        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))
        except Exception as e:
            _log.warning("session.persist_failed", error=str(e))

    async def load_from_disk(self) -> int:
        """
        Load sessions from disk and validate that their kernels are still alive.

        Stale sessions (dead kernels) are discarded.  If the file is missing,
        corrupted, or has a version mismatch, it is silently ignored.

        Returns:
            Number of sessions successfully restored.
        """
        if self._persistence_path is None or not self._persistence_path.exists():
            return 0

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

        if not isinstance(raw, dict) or raw.get("version") != self._PERSISTENCE_VERSION:
            logger.warning(
                f"Sessions file has unexpected version "
                f"(expected {self._PERSISTENCE_VERSION}, got {raw.get('version') if isinstance(raw, dict) else '?'}), discarding"
            )
            return 0

        saved_sessions = raw.get("sessions", {})
        restored = 0

        for path, s in saved_sessions.items():
            kernel_id = s.get("kernel_id")
            if not kernel_id:
                continue

            # Validate kernel is still alive
            conn_info = self._kernel_backend.get_kernel_connection_info(kernel_id)
            if conn_info is None:
                logger.info(
                    f"Discarding stale session for {path}: "
                    f"kernel {kernel_id} no longer exists"
                )
                continue

            try:
                session = NotebookSession(
                    session_id=s["session_id"],
                    kernel_id=kernel_id,
                    notebook_path=s["notebook_path"],
                    created_at=datetime.fromisoformat(s["created_at"]),
                    kernel_name=s.get("kernel_name", "python3"),
                    python_path=s.get("python_path"),
                    env_name=s.get("env_name"),
                )
                self._sessions[path] = session
                restored += 1
                logger.info(
                    f"Restored session {session.session_id} for {path} "
                    f"(kernel {kernel_id})"
                )
            except (KeyError, ValueError) as e:
                logger.warning(f"Skipping malformed session entry for {path}: {e}")

        # Save cleaned state (only valid sessions)
        self._save_to_disk()
        return restored

    async def get_or_create_session(
        self,
        notebook_path: str,
        kernel_name: str = "python3",
        python_path: Optional[str] = None,
        env_name: Optional[str] = None,
    ) -> Tuple[NotebookSession, bool]:
        """
        Get existing session for notebook or create a new one.

        Args:
            notebook_path: Path to the notebook file (relative to project root).
            kernel_name: Jupyter kernel spec name (default: python3).
            python_path: Optional path to Python executable for custom environments.
            env_name: Optional display name for the environment.

        Returns:
            Tuple of (NotebookSession, was_created) where was_created is True
            if a new session was created, False if an existing session was returned.

        Raises:
            RuntimeError: If kernel fails to start.
        """
        normalized_path = self._normalize_path(notebook_path)

        async with self._async_lock:
            # Return existing session if kernel is still alive
            if normalized_path in self._sessions:
                session = self._sessions[normalized_path]

                # Validate kernel still exists (may have died since session was created)
                conn_info = self._kernel_backend.get_kernel_connection_info(
                    session.kernel_id
                )
                if conn_info is not None:
                    logger.debug(
                        f"Returning existing session {session.session_id} "
                        f"for notebook {normalized_path}"
                    )
                    return session, False

                # Kernel is dead — remove stale session and create a new one
                logger.warning(
                    f"Stale session {session.session_id} for {normalized_path}: "
                    f"kernel {session.kernel_id} no longer exists, creating new session"
                )
                del self._sessions[normalized_path]

            # Create new kernel/session via KernelBackend
            logger.info(f"Creating new session for notebook: {normalized_path}")

            try:
                kernel_id = await self._kernel_backend.start_kernel(
                    kernel_name=kernel_name,
                    python_path=python_path,
                    env_name=env_name,
                )

                # Get session ID from kernel connection info
                conn_info = self._kernel_backend.get_kernel_connection_info(kernel_id)
                if conn_info is None:
                    # Cleanup the orphaned kernel before raising
                    try:
                        await self._kernel_backend.stop_kernel(kernel_id)
                    except Exception:
                        logger.warning(f"Failed to cleanup orphaned kernel {kernel_id}")
                    raise RuntimeError(f"Kernel {kernel_id} started but connection not found")

                session = NotebookSession(
                    session_id=conn_info["session_id"],
                    kernel_id=kernel_id,
                    notebook_path=normalized_path,
                    created_at=datetime.now(),
                    kernel_name=kernel_name,
                    python_path=python_path,
                    env_name=env_name,
                )

                self._sessions[normalized_path] = session
                self._save_to_disk()

                _log.info(
                    "session.create",
                    session_id=session.session_id,
                    kernel_id=kernel_id,
                    notebook_path=normalized_path,
                )

                return session, True

            except Exception as e:
                logger.error(f"Failed to create session for {normalized_path}: {e}")
                raise RuntimeError(f"Failed to create session: {e}") from e

    async def close_session(self, notebook_path: str) -> bool:
        """
        Close a session and stop its kernel if no other sessions are using it.

        If the kernel is shared with other notebooks, the session is removed
        but the kernel is kept alive for the other notebooks.

        Args:
            notebook_path: Path to the notebook file.

        Returns:
            True if session was closed, False if no session existed.
        """
        normalized_path = self._normalize_path(notebook_path)

        async with self._async_lock:
            session = self._sessions.pop(normalized_path, None)

            if session is None:
                logger.debug(f"No session found for notebook: {normalized_path}")
                return False

            self._save_to_disk()

            # Check if other notebooks are using this kernel (kernel sharing)
            kernel_users = [
                s.notebook_path
                for s in self._sessions.values()
                if s.kernel_id == session.kernel_id
            ]

            if kernel_users:
                # Kernel is shared - don't stop it
                logger.info(
                    f"Session for {normalized_path} closed. "
                    f"Kernel {session.kernel_id} kept alive for: {kernel_users}"
                )
                return True

            # No other users, safe to stop the kernel
            logger.info(
                f"Closing session {session.session_id} and stopping kernel "
                f"for notebook {normalized_path}"
            )

            try:
                await self._kernel_backend.stop_kernel(session.kernel_id)
                logger.info(f"Session {session.session_id} closed successfully")
            except Exception as e:
                logger.warning(
                    f"Error stopping kernel for session {session.session_id}: {e}"
                )
                # Session is already removed from our tracking, so consider it closed

            return True

    def get_session_for_path(self, notebook_path: str) -> Optional[NotebookSession]:
        """
        Get the session for a notebook path, if one exists.

        Args:
            notebook_path: Path to the notebook file.

        Returns:
            NotebookSession if found, None otherwise.
        """
        normalized_path = self._normalize_path(notebook_path)
        return self._sessions.get(normalized_path)

    def get_kernel_for_path(self, notebook_path: str) -> Optional[str]:
        """
        Get the kernel ID for a notebook path, if one exists.

        Args:
            notebook_path: Path to the notebook file.

        Returns:
            Kernel ID if found, None otherwise.
        """
        session = self.get_session_for_path(notebook_path)
        return session.kernel_id if session else None

    def list_sessions(self) -> List[Dict]:
        """
        List all active sessions.

        Returns:
            List of session info dicts with keys:
            - session_id: Jupyter session ID
            - kernel_id: Kernel ID
            - notebook_path: Path to notebook
            - created_at: ISO timestamp of session creation
            - kernel_name: Kernel spec name
            - python_path: Custom Python path (if any)
            - env_name: Environment display name (if any)
        """
        from skop_runner.file_manager import FileManager
        normalize = FileManager._normalize_output_path
        return [
            {
                "session_id": session.session_id,
                "kernel_id": session.kernel_id,
                "notebook_path": normalize(session.notebook_path),
                "created_at": session.created_at.isoformat(),
                "kernel_name": session.kernel_name,
                "python_path": session.python_path,
                "env_name": session.env_name,
            }
            for session in self._sessions.values()
        ]

    async def list_sessions_with_status(self) -> List[Dict]:
        """
        List all active sessions enriched with live kernel execution state.

        Queries get_kernel_status() per session to include the real-time
        execution state (idle/busy/starting). Used by state.snapshot RPC.

        Returns:
            List of session info dicts (same as list_sessions) plus:
            - execution_state: Live kernel state ("idle", "busy", "starting", "unknown")
        """
        sessions = self.list_sessions()

        for session_dict in sessions:
            try:
                status_result = await self._kernel_backend.get_kernel_status(
                    session_dict["kernel_id"]
                )
                if status_result.get("status") == "not_found":
                    session_dict["execution_state"] = "dead"
                else:
                    session_dict["execution_state"] = status_result.get(
                        "execution_state", "unknown"
                    )
            except Exception as e:
                logger.warning(
                    f"Failed to get status for kernel {session_dict['kernel_id']}: {e}"
                )
                session_dict["execution_state"] = "unknown"

        return sessions

    async def close_all_sessions(self) -> Dict[str, List]:
        """
        Close all active sessions.

        Returns:
            Dict with 'closed' (list of notebook paths) and 'failed' (list of paths with errors).
        """
        closed = []
        failed = []

        # Get all paths while holding lock, then release for individual closes
        async with self._async_lock:
            paths = list(self._sessions.keys())

        for path in paths:
            try:
                success = await self.close_session(path)
                if success:
                    closed.append(path)
            except Exception as e:
                logger.error(f"Error closing session for {path}: {e}")
                failed.append(path)

        return {"closed": closed, "failed": failed}

    def get_session_for_kernel(self, kernel_id: str) -> Optional[NotebookSession]:
        """
        Get the session that owns a kernel.

        Args:
            kernel_id: The kernel ID to look up.

        Returns:
            NotebookSession if found, None otherwise.
        """
        for session in self._sessions.values():
            if session.kernel_id == kernel_id:
                return session
        return None

    def remove_sessions_for_kernel(self, kernel_id: str) -> List[str]:
        """
        Remove all sessions associated with a dead kernel.

        Called when a kernel dies to prevent stale sessions from being
        returned by get_or_create_session.

        Args:
            kernel_id: The dead kernel's ID.

        Returns:
            List of notebook paths whose sessions were removed.
        """
        removed: List[str] = []
        for path, session in list(self._sessions.items()):
            if session.kernel_id == kernel_id:
                del self._sessions[path]
                removed.append(path)
                logger.info(
                    f"Removed stale session {session.session_id} "
                    f"for notebook {path} (kernel {kernel_id} died)"
                )
        if removed:
            self._save_to_disk()
        return removed

    def rename_session(
        self, old_path: str, new_path: str, is_folder: bool = False
    ) -> bool:
        """
        Re-key sessions after a file or folder rename.

        For a single file rename, pops the session at old_path and re-inserts
        it at new_path with an updated notebook_path field.  For a folder
        rename, applies prefix replacement to every session whose path starts
        with ``old_path + "/"``.

        Args:
            old_path: The previous notebook (or folder) path.
            new_path: The new notebook (or folder) path.
            is_folder: When True, treat old_path as a folder prefix.

        Returns:
            True if at least one session was renamed, False otherwise.
        """
        old_normalized = self._normalize_path(old_path)
        new_normalized = self._normalize_path(new_path)

        if old_normalized == new_normalized:
            return False

        renamed = False

        if is_folder:
            # Folder rename: prefix replacement on all matching sessions
            prefix = old_normalized + os.sep
            to_rename = [
                (p, s)
                for p, s in self._sessions.items()
                if p == old_normalized or p.startswith(prefix)
            ]
            for path, session in to_rename:
                del self._sessions[path]
                if path == old_normalized:
                    updated = new_normalized
                else:
                    updated = new_normalized + path[len(old_normalized) :]
                session.notebook_path = updated
                self._sessions[updated] = session
                renamed = True
                logger.info(
                    f"Renamed session {session.session_id}: {path} -> {updated}"
                )
        else:
            # Single file rename
            session = self._sessions.pop(old_normalized, None)
            if session is not None:
                session.notebook_path = new_normalized
                self._sessions[new_normalized] = session
                renamed = True
                logger.info(
                    f"Renamed session {session.session_id}: "
                    f"{old_normalized} -> {new_normalized}"
                )

        if renamed:
            self._save_to_disk()

        return renamed

    async def connect_to_kernel(
        self,
        notebook_path: str,
        kernel_id: str,
    ) -> NotebookSession:
        """
        Connect a notebook to an existing kernel (kernel sharing).

        This creates a session entry without starting a new kernel,
        allowing multiple notebooks to share the same kernel state.

        Args:
            notebook_path: Path to the notebook file.
            kernel_id: ID of existing kernel to connect to.

        Returns:
            NotebookSession linked to the shared kernel.

        Raises:
            ValueError: If kernel doesn't exist.
            RuntimeError: If notebook already has a session.
        """
        normalized_path = self._normalize_path(notebook_path)

        async with self._async_lock:
            # Check if notebook already has a session
            if normalized_path in self._sessions:
                raise RuntimeError(
                    f"Notebook {normalized_path} already has a session. "
                    "Close it first before connecting to a different kernel."
                )

            # Verify kernel exists
            conn_info = self._kernel_backend.get_kernel_connection_info(kernel_id)
            if conn_info is None:
                raise ValueError(f"Kernel {kernel_id} not found")

            # Find the original session for this kernel
            original_session = self.get_session_for_kernel(kernel_id)

            session = NotebookSession(
                session_id=conn_info["session_id"],
                kernel_id=kernel_id,
                notebook_path=normalized_path,
                created_at=datetime.now(),
                kernel_name=original_session.kernel_name if original_session else "python3",
                python_path=conn_info["python_path"],
                env_name=conn_info["env_name"],
            )

            self._sessions[normalized_path] = session
            self._save_to_disk()

            logger.info(
                f"Connected notebook {normalized_path} to existing kernel {kernel_id}"
            )

            return session

    async def disconnect_from_kernel(self, notebook_path: str) -> bool:
        """
        Disconnect a notebook from a shared kernel without stopping the kernel.

        Use this instead of close_session when a notebook is sharing a kernel
        and you want to disconnect without killing the kernel for other notebooks.

        Args:
            notebook_path: Path to the notebook file.

        Returns:
            True if disconnected, False if no session existed.
        """
        normalized_path = self._normalize_path(notebook_path)

        async with self._async_lock:
            session = self._sessions.pop(normalized_path, None)

            if session is None:
                logger.debug(f"No session found for notebook: {normalized_path}")
                return False

            self._save_to_disk()

            # Check if other notebooks are using this kernel
            kernel_users = [
                s.notebook_path
                for s in self._sessions.values()
                if s.kernel_id == session.kernel_id
            ]

            if kernel_users:
                logger.info(
                    f"Disconnected {normalized_path} from kernel {session.kernel_id}. "
                    f"Kernel still in use by: {kernel_users}"
                )
            else:
                logger.info(
                    f"Disconnected {normalized_path} from kernel {session.kernel_id}. "
                    "No other users - kernel may be stopped by close_session if desired."
                )

            return True
