#!/usr/bin/env python
# coding: utf-8

# # Gett
# 
# 
# ## Step 1: list kernels and sessions
# 
# In the Jupyter server API, a `session` is a relationship between a document (or console) and a kernel.
# In practice, it is typical for each kernel to be associated with one document.
# 
# **Caveats**
# 
# - Kernels do not _need_ to be associated with a document,
#   so listing `kernels` directly may discover more kernels than the sessions API,
#   but this is rare.
# - Technically, one kernel can be associated with multiple documents,
#   but this is also exceedingly rare and not super relevant.

# In[1]:


import requests

token = "abc123"  # get this somewhere, e.g. $JUPYTERHUB_API_TOKEN or via a JupyterHub service token

s = requests.Session()
s.headers = {"Authorization": "token abc123"}
server_url = "http://127.0.0.1:8888"


# In[2]:


session_list = s.get(f"{server_url}/api/sessions").json()
session_list


# In[3]:


kernel_list = s.get(f"{server_url}/api/kernels").json()
kernel_list


# ## Step 2: look up kernel PIDs
# 
# Jupyter doesn't include PIDs in any part of the API because kernels need not be local to the Jupyter server
# (e.g. when they are run via a gateway server).
# However, they are often enough that the IPython kernel has a prototype message handler called a `usage_request`.
# We can use this to ask for information, including the PID of the process.
# 
# **Note:** this step only works for recent versions of `ipykernel`. This is a nonstandard message, and not supported by other kernels.
# An alternative would be to use an `execute_request` to run `os.getpid()`. This, in turn, is Python-specific and would need to be implemented.
# 
# 
# To do this, we connect a websocket to the Jupyter server to send messages directly to the kernel:

# In[7]:


import json

from tornado.websocket import websocket_connect
from jupyter_client.session import Session, json_packer


# In[9]:


client_session = Session()

kernel_id = session_list[0]["kernel"]["id"]

ws = await websocket_connect(f"ws{server_url[4:]}/api/kernels/{kernel_id}/channels?token={token}")
msg = client_session.msg("usage_request", content={})
msg["channel"] = "control"
json_msg = json_packer(msg)
await ws.write_message(json_msg)
while True:
    reply = json.loads(await ws.read_message())
    if reply["channel"] == "control":
        break
reply["content"]


# We can use this information to produce a map of PID to notebook document / console

# In[10]:


async def get_kernel_pid(kernel_id):
    client_session = Session()
    ws = await websocket_connect(f"ws{server_url[4:]}/api/kernels/{kernel_id}/channels?token={token}")
    msg = client_session.msg("usage_request", content={})
    msg["channel"] = "control"
    json_msg = json_packer(msg)
    await ws.write_message(json_msg)
    while True:
        reply = json.loads(await ws.read_message())
        if reply["channel"] == "control":
            break
    ws.close()
    return reply["content"]["pid"]

for session in session_list:
    pid = await get_kernel_pid(session["kernel"]["id"])
    print(f"{pid}: {session['type']} {session['name']!r}")


# ## Alternative: using kernel ids from connection files
# 
# In [the question](https://discourse.jupyter.org/t/how-i-could-get-the-notebook-name-having-the-kernel-json-runtime-key-or-os-process-id/15168/2),
# the runtime file is also mentioned as a valid input.
# This may be simpler and more general to use.
# 
# By convention, Jupyter servers put the kernel id in the connection file as well. So a simpler approach that works for all kernel languages
# is to check for the kernel id in.
# 
# You could know _exactly_ how we produce these files and parse them out, but perhaps simpler and more robust to internal detail changes is simple substring matching:

# In[13]:


from jupyter_core.paths import jupyter_runtime_dir
import glob


# In[14]:


connection_files = glob.glob(f"{jupyter_runtime_dir()}/kernel-*.json")
connection_files


# In[16]:


import os

for session in session_list:
    kernel_id = session["kernel"]["id"]
    for connection_file in connection_files:
        fname = os.path.basename(connection_file)
        if kernel_id in fname:
            print(f"{fname}: {session['type']} {session['name']!r}")
            break