"""
Jupyter Server REST API Client.

Provides async HTTP client for Jupyter Server REST API with:
- Kernel lifecycle (start, stop, restart, interrupt)
- Session management (create, list, get, delete)
- Kernelspec discovery
- Health checks
"""

import asyncio
import logging
import re
from typing import Any, Dict, List, Optional, Union
from urllib.parse import urlencode, urljoin

import httpx

logger = logging.getLogger("skop-runner")

# UUID pattern for validating kernel/session IDs (security: prevents path injection)
_UUID_PATTERN = re.compile(
    r"^[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}$"
)

# Timeouts (seconds)
DEFAULT_TIMEOUT = 30.0
CONNECT_TIMEOUT = 10.0

# Type alias for JSON responses
JSONValue = Union[Dict[str, Any], List[Any], str, int, float, bool, None]


# =============================================================================
# Exceptions
# =============================================================================


class JupyterAPIError(Exception):
    """Base exception for Jupyter API errors."""

    def __init__(self, message: str, status_code: Optional[int] = None):
        super().__init__(message)
        self.status_code = status_code


class JupyterConnectionError(JupyterAPIError):
    """Network or connection error (timeout, refused, etc.)."""

    def __init__(self, message: str):
        super().__init__(message, status_code=None)


class JupyterAuthenticationError(JupyterAPIError):
    """Authentication failed (401, 403)."""

    pass


class JupyterNotFoundError(JupyterAPIError):
    """Resource not found (404)."""

    pass


class JupyterServerError(JupyterAPIError):
    """Server-side error (500+)."""

    pass


class JupyterConflictError(JupyterAPIError):
    """Resource conflict (409). Common during concurrent session operations."""

    pass


class JupyterGoneError(JupyterAPIError):
    """Resource is gone (410). Kernel died before session cleanup."""

    pass


# =============================================================================
# API Client
# =============================================================================


class JupyterAPIClient:
    """
    HTTP client for Jupyter Server REST API.

    Thread-safe for concurrent use from multiple coroutines.
    Must call close() when done to release HTTP connections.

    Example:
        async with JupyterAPIClient("http://127.0.0.1:18888", "token") as client:
            kernels = await client.list_kernels()
    """

    def __init__(
        self,
        base_url: str,
        token: str,
        timeout: float = DEFAULT_TIMEOUT,
    ):
        """
        Initialize the API client.

        Args:
            base_url: Jupyter Server base URL (e.g., "http://127.0.0.1:18888").
            token: Authentication token for Jupyter Server.
            timeout: Default request timeout in seconds.

        Raises:
            ValueError: If base_url is invalid.
        """
        self._base_url = self._sanitize_base_url(base_url)
        self._token = token
        self._timeout = timeout
        self._client: Optional[httpx.AsyncClient] = None
        # Lazy-initialized lock protects _client, _base_url, _token
        self.__lock: Optional[asyncio.Lock] = None

    @property
    def _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

    @property
    def base_url(self) -> str:
        """Get the base URL (without trailing slash)."""
        return self._base_url

    @property
    def ws_base_url(self) -> str:
        """Get the WebSocket base URL (ws:// or wss://)."""
        return self._base_url.replace("http://", "ws://").replace("https://", "wss://")

    async def _get_client(self) -> httpx.AsyncClient:
        """Get or create the HTTP client (lazy initialization with double-check locking)."""
        if self._client is not None:
            return self._client

        async with self._lock:
            if self._client is None:
                self._client = httpx.AsyncClient(
                    headers={"Authorization": f"token {self._token}"},
                    timeout=httpx.Timeout(self._timeout, connect=CONNECT_TIMEOUT),
                )
            return self._client

    async def close(self) -> None:
        """Close the HTTP client and release connections. Idempotent and thread-safe."""
        async with self._lock:
            client = self._client
            self._client = None

        # Close outside the lock to avoid holding it during I/O
        if client is not None:
            await client.aclose()

    async def update_credentials(self, base_url: str, token: str) -> None:
        """
        Update credentials after server restart. Closes existing client.

        Thread-safe: acquires lock to prevent races with _get_client().
        New client will be created on next request with new credentials.

        Args:
            base_url: New Jupyter Server base URL.
            token: New authentication token.
        """
        # Validate before acquiring lock to fail fast
        normalized_url = self._sanitize_base_url(base_url)

        async with self._lock:
            client = self._client
            self._client = None
            self._base_url = normalized_url
            self._token = token

        # Close outside the lock to avoid holding it during I/O
        if client is not None:
            await client.aclose()

    async def __aenter__(self) -> "JupyterAPIClient":
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        await self.close()

    def _validate_id(self, id_value: str, name: str = "id") -> None:
        """Validate that an ID is a valid UUID (prevents path injection)."""
        if not _UUID_PATTERN.match(id_value):
            raise ValueError(f"Invalid {name}: must be a valid UUID, got '{id_value}'")

    def _sanitize_base_url(self, base_url: str) -> str:
        """Ensure base URL uses HTTP(S) and trim trailing slash."""
        if not base_url.startswith(("http://", "https://")):
            raise ValueError("base_url must start with http:// or https://")
        return base_url.rstrip("/")

    def _build_url(self, path: str) -> str:
        """Build full URL preserving any base URL prefix."""
        base = self._base_url if self._base_url.endswith("/") else self._base_url + "/"
        return urljoin(base, path.lstrip("/"))

    def kernel_ws_url(self, kernel_id: str) -> str:
        """
        Get authenticated WebSocket URL for kernel channels.

        Args:
            kernel_id: The kernel UUID.

        Returns:
            WebSocket URL with token for authentication.
        """
        self._validate_id(kernel_id, "kernel_id")
        base = self.ws_base_url if self.ws_base_url.endswith("/") else self.ws_base_url + "/"
        url = urljoin(base, f"api/kernels/{kernel_id}/channels")
        return f"{url}?{urlencode({'token': self._token})}"

    async def _request(
        self,
        method: str,
        path: str,
        json: Optional[Dict[str, Any]] = None,
        timeout: Optional[float] = None,
    ) -> JSONValue:
        """
        Make an HTTP request to the Jupyter Server API.

        Args:
            method: HTTP method (GET, POST, DELETE, PATCH).
            path: API path (e.g., "/api/kernels").
            json: Request body as dict.
            timeout: Override default timeout for this request.

        Returns:
            Parsed JSON response, or None for 204 No Content.

        Raises:
            JupyterConnectionError: Network/timeout error.
            JupyterAuthenticationError: 401/403 response.
            JupyterNotFoundError: 404 response.
            JupyterServerError: 500+ response.
            JupyterAPIError: Other HTTP errors.
        """
        url = self._build_url(path)
        logger.debug(f"API: {method} {url}")

        client = await self._get_client()

        # Build kwargs (don't pass timeout=None, it disables timeouts in httpx)
        kwargs: Dict[str, Any] = {"url": url}
        if json is not None:
            kwargs["json"] = json
        if timeout is not None:
            kwargs["timeout"] = timeout

        try:
            response = await client.request(method, **kwargs)
        except httpx.TimeoutException as e:
            raise JupyterConnectionError(f"Request timed out: {path}") from e
        except httpx.ConnectError as e:
            raise JupyterConnectionError(
                f"Failed to connect to Jupyter Server at {self._base_url}"
            ) from e
        except httpx.RequestError as e:
            raise JupyterConnectionError(f"Request failed: {e}") from e

        return self._handle_response(response, path)

    def _handle_response(self, response: httpx.Response, path: str) -> JSONValue:
        """Handle HTTP response and translate errors to exceptions."""
        status = response.status_code

        if status == 204:
            return None

        if 200 <= status < 300:
            try:
                return response.json()
            except Exception as e:
                raise JupyterAPIError(
                    f"Invalid JSON response from {path}", status_code=status
                ) from e

        # Extract error message from response
        try:
            data = response.json()
            error_msg = data.get("message", data.get("reason", str(data)))
        except Exception:
            error_msg = response.text[:200] if response.text else "Unknown error"

        if status in (401, 403):
            raise JupyterAuthenticationError(
                f"Authentication failed for {path}: {error_msg}", status_code=status
            )
        if status == 404:
            raise JupyterNotFoundError(
                f"Resource not found: {path}: {error_msg}", status_code=status
            )
        if status == 409:
            raise JupyterConflictError(
                f"Resource conflict for {path}: {error_msg}", status_code=status
            )
        if status == 410:
            raise JupyterGoneError(
                f"Resource gone: {path}: {error_msg}", status_code=status
            )
        if status >= 500:
            raise JupyterServerError(
                f"Server error ({status}) for {path}: {error_msg}", status_code=status
            )

        raise JupyterAPIError(
            f"API error ({status}) for {path}: {error_msg}", status_code=status
        )

    # =========================================================================
    # Kernel API Methods
    # =========================================================================

    async def list_kernels(self) -> List[Dict[str, Any]]:
        """
        List all running kernels.

        Returns:
            List of kernel info dicts with keys: id, name, execution_state, etc.
        """
        result = await self._request("GET", "/api/kernels")
        return result if isinstance(result, list) else []

    async def start_kernel(
        self, name: str = "python3", path: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Start a new kernel.

        Args:
            name: Kernel spec name (e.g., "python3").
            path: Optional working directory path for the kernel.

        Returns:
            Kernel info dict with keys: id, name, execution_state, etc.
        """
        body: Dict[str, Any] = {"name": name}
        if path is not None:
            body["path"] = path
        result = await self._request("POST", "/api/kernels", json=body)
        return result  # type: ignore[return-value]

    async def get_kernel(self, kernel_id: str) -> Dict[str, Any]:
        """
        Get information about a specific kernel.

        Args:
            kernel_id: The kernel UUID.

        Returns:
            Kernel info dict.

        Raises:
            JupyterNotFoundError: If kernel doesn't exist.
        """
        self._validate_id(kernel_id, "kernel_id")
        result = await self._request("GET", f"/api/kernels/{kernel_id}")
        return result  # type: ignore[return-value]

    async def delete_kernel(self, kernel_id: str) -> None:
        """
        Terminate a kernel.

        Args:
            kernel_id: The kernel UUID.

        Raises:
            JupyterNotFoundError: If kernel doesn't exist.
        """
        self._validate_id(kernel_id, "kernel_id")
        await self._request("DELETE", f"/api/kernels/{kernel_id}")

    async def interrupt_kernel(self, kernel_id: str) -> None:
        """
        Send an interrupt signal to a kernel.

        Args:
            kernel_id: The kernel UUID.

        Raises:
            JupyterNotFoundError: If kernel doesn't exist.
        """
        self._validate_id(kernel_id, "kernel_id")
        await self._request("POST", f"/api/kernels/{kernel_id}/interrupt")

    async def restart_kernel(self, kernel_id: str) -> Dict[str, Any]:
        """
        Restart a kernel.

        The kernel will be stopped and started again. Existing WebSocket
        connections will be invalidated.

        Args:
            kernel_id: The kernel UUID.

        Returns:
            Updated kernel info dict.

        Raises:
            JupyterNotFoundError: If kernel doesn't exist.
        """
        self._validate_id(kernel_id, "kernel_id")
        result = await self._request("POST", f"/api/kernels/{kernel_id}/restart")
        return result  # type: ignore[return-value]

    # =========================================================================
    # Session API Methods
    # =========================================================================

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

        Returns:
            List of session info dicts with keys: id, path, kernel, etc.
        """
        result = await self._request("GET", "/api/sessions")
        return result if isinstance(result, list) else []

    async def create_session(
        self,
        path: str,
        kernel_name: str = "python3",
        name: Optional[str] = None,
        session_type: str = "notebook",
    ) -> Dict[str, Any]:
        """
        Create a new session (which also starts a kernel).

        Args:
            path: Path to the notebook or file.
            kernel_name: Kernel spec name (e.g., "python3").
            name: Display name for the session (defaults to filename from path).
            session_type: Session type ("notebook", "console", "file").

        Returns:
            Session info dict with keys: id, path, name, kernel, etc.
            The kernel dict contains: id, name, execution_state.
        """
        # Default name to filename from path
        if name is None:
            name = path.rsplit("/", 1)[-1] if "/" in path else path

        body = {
            "path": path,
            "name": name,
            "type": session_type,
            "kernel": {"name": kernel_name},
        }
        result = await self._request("POST", "/api/sessions", json=body)
        return result  # type: ignore[return-value]

    async def get_session(self, session_id: str) -> Dict[str, Any]:
        """
        Get information about a specific session.

        Args:
            session_id: The session UUID.

        Returns:
            Session info dict.

        Raises:
            JupyterNotFoundError: If session doesn't exist.
        """
        self._validate_id(session_id, "session_id")
        result = await self._request("GET", f"/api/sessions/{session_id}")
        return result  # type: ignore[return-value]

    async def delete_session(self, session_id: str) -> None:
        """
        Delete a session (which also terminates its kernel).

        Args:
            session_id: The session UUID.

        Raises:
            JupyterNotFoundError: If session doesn't exist.
            JupyterGoneError: If the kernel died before session cleanup.
        """
        self._validate_id(session_id, "session_id")
        await self._request("DELETE", f"/api/sessions/{session_id}")

    async def update_session(
        self,
        session_id: str,
        path: Optional[str] = None,
        name: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Update a session's path or name.

        Args:
            session_id: The session UUID.
            path: New path for the session.
            name: New display name for the session.

        Returns:
            Updated session info dict.

        Raises:
            ValueError: If neither path nor name is provided.
            JupyterNotFoundError: If session doesn't exist.
        """
        if path is None and name is None:
            raise ValueError("At least one of path or name must be provided")

        self._validate_id(session_id, "session_id")

        body: Dict[str, str] = {}
        if path is not None:
            body["path"] = path
        if name is not None:
            body["name"] = name

        result = await self._request("PATCH", f"/api/sessions/{session_id}", json=body)
        return result  # type: ignore[return-value]

    # =========================================================================
    # Kernelspec API Methods
    # =========================================================================

    async def get_kernelspecs(self) -> Dict[str, Any]:
        """
        Get all available kernel specifications.

        Returns:
            Dict with keys: default (str), kernelspecs (dict of name -> spec).
        """
        result = await self._request("GET", "/api/kernelspecs")
        return result  # type: ignore[return-value]

    async def get_kernelspec(self, name: str) -> Dict[str, Any]:
        """
        Get a specific kernel specification.

        Args:
            name: Kernel spec name (e.g., "python3").

        Returns:
            Kernelspec dict with keys: name, spec, resources.

        Raises:
            JupyterNotFoundError: If kernelspec doesn't exist.
        """
        result = await self._request("GET", f"/api/kernelspecs/{name}")
        return result  # type: ignore[return-value]

    # =========================================================================
    # Status API Methods
    # =========================================================================

    async def get_status(self) -> Dict[str, Any]:
        """
        Get server status information.

        Returns:
            Status dict with keys like: started, last_activity, connections, kernels.
        """
        result = await self._request("GET", "/api/status")
        return result  # type: ignore[return-value]

    async def is_healthy(self, timeout: float = 5.0) -> bool:
        """
        Check if the server is healthy and responding.

        This is a convenience method for health checks. It catches all
        exceptions and returns False instead of raising.

        Args:
            timeout: Request timeout in seconds (default 5s for quick checks).

        Returns:
            True if server responds successfully, False otherwise.
        """
        try:
            await self._request("GET", "/api/status", timeout=timeout)
            return True
        except JupyterAPIError:
            return False
