fn_registry/python/functions/notebook/jupyter_discover.py

"""Descubrimiento de instancias Jupyter Lab activas via API REST."""

import json
import os
import urllib.error
import urllib.request
from pathlib import Path


_DEFAULT_PORTS = [8888, 8889, 8890, 8891, 8892]


def _get(url: str, timeout: float = 2.0) -> dict | list | None:
    """Hace GET a url y retorna el JSON parseado, o None si falla."""
    try:
        with urllib.request.urlopen(url, timeout=timeout) as resp:
            return json.loads(resp.read().decode())
    except Exception:
        return None


def _is_collaborative(config: dict | list | None) -> bool:
    """Detecta si el servidor tiene jupyter-collaboration/YDocExtension activo."""
    if not isinstance(config, dict):
        return False
    # Jupyter Lab expone la config de extensiones bajo claves como
    # 'LabApp' o similares; la presencia de 'collaborative' o 'YDocExtension'
    # en cualquier valor de primer nivel indica modo colaborativo.
    raw = json.dumps(config).lower()
    return "ydocextension" in raw or "collaborative" in raw


def _query_instance(base_url: str) -> dict | None:
    """Consulta la API REST de una instancia Jupyter y retorna su estado.

    Retorna None si la instancia no responde o no es Jupyter.
    """
    status = _get(f"{base_url}/api/status")
    if status is None:
        return None

    config = _get(f"{base_url}/api/config")
    kernels_raw = _get(f"{base_url}/api/kernels") or []
    sessions_raw = _get(f"{base_url}/api/sessions") or []

    kernels = []
    if isinstance(kernels_raw, list):
        for k in kernels_raw:
            if isinstance(k, dict):
                kernels.append({
                    "id": k.get("id", ""),
                    "name": k.get("name", ""),
                    "execution_state": k.get("execution_state", ""),
                    "last_activity": k.get("last_activity", ""),
                })

    sessions = []
    if isinstance(sessions_raw, list):
        for s in sessions_raw:
            if isinstance(s, dict):
                kernel = s.get("kernel") or {}
                path = s.get("path") or s.get("notebook", {}).get("path", "")
                sessions.append({
                    "notebook": path,
                    "kernel_id": kernel.get("id", ""),
                    "kernel_state": kernel.get("execution_state", ""),
                })

    return {
        "kernels": kernels,
        "sessions": sessions,
        "collaborative": _is_collaborative(config),
    }


def _scan_analysis_ports(registry_root: str) -> list[tuple[int, str]]:
    """Escanea subdirectorios de analysis/ buscando archivos .jupyter-port.

    Retorna lista de (puerto, nombre_analisis).
    """
    root = Path(registry_root) if registry_root else Path.cwd()
    analysis_dir = root / "analysis"
    results: list[tuple[int, str]] = []

    if not analysis_dir.is_dir():
        return results

    for entry in analysis_dir.iterdir():
        if not entry.is_dir():
            continue
        port_file = entry / ".jupyter-port"
        if port_file.is_file():
            try:
                port = int(port_file.read_text().strip())
                results.append((port, entry.name))
            except (ValueError, OSError):
                pass

    return results


def jupyter_discover(
    registry_root: str = "",
    ports: list[int] | None = None,
) -> list[dict]:
    """Descubre instancias de Jupyter Lab activas consultando su API REST.

    Escanea primero los archivos .jupyter-port en subdirectorios de analysis/
    para encontrar puertos registrados, y luego aplica un fallback sobre puertos
    comunes (8888-8892). Para cada instancia que responde consulta /api/status,
    /api/config, /api/kernels y /api/sessions.

    Args:
        registry_root: Raiz del fn_registry. Si vacio usa el directorio actual
                       o la variable de entorno FN_REGISTRY_ROOT.
        ports: Lista de puertos a escanear. Si None, usa los puertos encontrados
               en .jupyter-port mas los defaults (8888-8892).

    Returns:
        Lista de dicts con: url, port, analysis, collaborative, kernels, sessions.
        Cada sesion incluye: notebook, kernel_id, kernel_state.
    """
    if not registry_root:
        registry_root = os.environ.get("FN_REGISTRY_ROOT", "")

    # Recopilar puertos a escanear
    port_analysis: dict[int, str] = {}

    if ports is not None:
        for p in ports:
            port_analysis[p] = ""
    else:
        # Primero los registrados en .jupyter-port
        for port, analysis_name in _scan_analysis_ports(registry_root):
            port_analysis[port] = analysis_name
        # Fallback: puertos comunes que no estén ya en la lista
        for p in _DEFAULT_PORTS:
            if p not in port_analysis:
                port_analysis[p] = ""

    results = []
    for port, analysis_name in port_analysis.items():
        base_url = f"http://localhost:{port}"
        info = _query_instance(base_url)
        if info is None:
            continue
        results.append({
            "url": base_url,
            "port": port,
            "analysis": analysis_name,
            "collaborative": info["collaborative"],
            "kernels": info["kernels"],
            "sessions": info["sessions"],
        })

    return results


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import argparse
    import sys

    parser = argparse.ArgumentParser(
        description="Descubre instancias de Jupyter Lab activas."
    )
    parser.add_argument(
        "--registry-root",
        default="",
        help="Raiz del fn_registry (default: FN_REGISTRY_ROOT env o cwd)",
    )
    parser.add_argument(
        "--port",
        dest="ports",
        type=int,
        action="append",
        metavar="PORT",
        help="Puerto a escanear (puede repetirse). Default: .jupyter-port + 8888-8892",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Emitir salida en JSON",
    )
    args = parser.parse_args()

    instances = jupyter_discover(
        registry_root=args.registry_root,
        ports=args.ports,
    )

    if args.json:
        print(json.dumps(instances, indent=2))
        sys.exit(0)

    if not instances:
        print("No se encontraron instancias de Jupyter Lab activas.")
        sys.exit(0)

    for inst in instances:
        label = f"  analysis: {inst['analysis']}" if inst["analysis"] else ""
        collab = "colaborativo" if inst["collaborative"] else "estandar"
        print(f"Jupyter Lab en {inst['url']}  [{collab}]{label}")
        if inst["kernels"]:
            print(f"  Kernels ({len(inst['kernels'])}):")
            for k in inst["kernels"]:
                print(f"    - {k['name']}  estado={k['execution_state']}  id={k['id'][:8]}...")
        else:
            print("  Kernels: ninguno")
        if inst["sessions"]:
            print(f"  Sesiones ({len(inst['sessions'])}):")
            for s in inst["sessions"]:
                print(f"    - {s['notebook']}  kernel={s['kernel_id'][:8]}...  estado={s['kernel_state']}")
        else:
            print("  Sesiones: ninguna")
        print()