"""JupyterLab notebook lifecycle."""
from __future__ import annotations

import time
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from .client import Client


class Notebook:
    """Wrapper around a notebook API response with lifecycle helpers."""

    def __init__(self, data: dict, client: Client):
        self._data = data
        self._client = client

    @property
    def id(self) -> str:
        return self._data.get("id", "")

    @property
    def name(self) -> str:
        return self._data.get("name", "")

    @property
    def status(self) -> str:
        return self._data.get("status", "")

    @property
    def jupyter_url(self) -> str:
        return self._data.get("jupyter_url", "")

    @property
    def gpu_type(self) -> str:
        return self._data.get("gpu_type", "")

    @property
    def created_at(self) -> str:
        return self._data.get("created_at", "")

    def refresh(self) -> None:
        """Re-fetch notebook state from the API."""
        self._data = self._client.get(f"/v1/research/notebooks/{self.id}")

    def wait_until_ready(
        self, timeout: float = 300, poll_interval: float = 5.0
    ) -> None:
        """Block until the notebook is running or has failed."""
        deadline = time.time() + timeout
        while time.time() < deadline:
            self.refresh()
            if self.status == "running":
                return
            if self.status == "failed":
                raise RuntimeError(f"Notebook {self.id} failed")
            time.sleep(min(poll_interval, max(0, deadline - time.time())))
        raise TimeoutError(
            f"Notebook {self.id} not ready within {timeout}s (status: {self.status})"
        )

    def stop(self) -> None:
        """Delete the notebook."""
        self._client.delete(f"/v1/research/notebooks/{self.id}")

    def __enter__(self) -> Notebook:
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.stop()


class Notebooks:
    def __init__(self, client: Client):
        self._client = client

    def create(
        self,
        name: str,
        *,
        gpu: str = "rtx-4090",
        gpu_count: int = 1,
        cpu: str = "4",
        memory: str = "8Gi",
        image: str = "",
        idle_timeout_minutes: int = 120,
        data_mounts: list[str] | None = None,
        env_vars: dict[str, str] | None = None,
        compute_target: str | None = None,
        provider: str | None = None,
        region: str | None = None,
        sovereignty: bool = False,
        min_vram: int | None = None,
    ) -> Notebook:
        """Create a JupyterLab notebook environment."""
        body: dict = {
            "name": name,
            "gpu_type": gpu,
            "gpu_count": gpu_count,
            "cpu": cpu,
            "memory": memory,
            "idle_timeout_minutes": idle_timeout_minutes,
        }
        if image:
            body["image"] = image
        if data_mounts:
            body["data_mounts"] = data_mounts
        if env_vars:
            body["env_vars"] = env_vars
        if compute_target:
            body["compute_target"] = compute_target
        if provider:
            body["provider"] = provider
        if region:
            body["region"] = region
        if sovereignty:
            body["sovereignty"] = True
        if min_vram:
            body["min_vram"] = min_vram
        data = self._client.post("/v1/research/notebooks", json=body)
        return Notebook(data, self._client)

    def get(self, notebook_id: str) -> Notebook:
        """Get notebook details."""
        data = self._client.get(f"/v1/research/notebooks/{notebook_id}")
        return Notebook(data, self._client)

    def delete(self, notebook_id: str) -> None:
        """Delete a notebook."""
        self._client.delete(f"/v1/research/notebooks/{notebook_id}")

    def list(self) -> list[dict]:
        """List all notebooks."""
        return self._client.get("/v1/research/notebooks")
