# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/routes/jupyter.ipynb.

# %% auto 0
__all__ = ['JupyterAPI_Error', 'JupyterAPI_WorkspaceStarted', 'get_jupyter_workspaces', 'get_jupyter_workspace_by_id',
           'parse_instance_service_location_and_prefix', 'get_workspace_auth_token_params', 'start_jupyter_workspace',
           'get_jupyter_content', 'generate_update_jupyter_body__new_content_path',
           'generate_update_jupyter_body__text', 'generate_update_jupyter_body__ipynb',
           'generate_update_jupyter_body__directory', 'generate_update_jupyter_body_factory',
           'generate_update_jupyter_body', 'create_jupyter_obj', 'delete_jupyter_content', 'update_jupyter_file',
           'get_content', 'update_jupyter_workspace_config']

# %% ../../nbs/routes/jupyter.ipynb 2
from domolibrary.client.DomoEntity import DomoEnum

import os
import httpx
import asyncio
from functools import partial

import urllib

import domolibrary.client.ResponseGetData as rgd
import domolibrary.client.DomoAuth as dmda
import domolibrary.client.get_data as gd

import domolibrary.client.DomoError as dmde
import domolibrary.utils.chunk_execution as dmce

# %% ../../nbs/routes/jupyter.ipynb 6
class JupyterAPI_Error(dmde.RouteError):
    def __init__(self, res, message: str = None):
        super().__init__(res=res, message=message or res.response)


class JupyterAPI_WorkspaceStarted(dmde.RouteError):
    def __init__(self, res, message: str = None):
        super().__init__(res=res, message=message or res.response)

# %% ../../nbs/routes/jupyter.ipynb 7
@gd.route_function
async def get_jupyter_workspaces(
    auth: dmda.DomoAuth,
    parent_class: str = None,
    session: httpx.AsyncClient = None,
    debug_num_stacks_to_drop=1,
    debug_api: bool = False,
    debug_loop: bool = False,
):
    url = f"https://{auth.domo_instance}.domo.com/api/datascience/v1/search/workspaces"

    body = {
        "limit": 50,
        "offset": 0,
        "sortFieldMap": {"CREATED": "DESC"},
        "filters": [],
    }

    def arr_fn(res):
        return res.response["workspaces"]

    offset_params = {"limit": "limit", "offset": "offset"}

    res = await gd.looper(
        url=url,
        method="POST",
        limit=50,
        body=body,
        auth=auth,
        arr_fn=arr_fn,
        offset_params_in_body=True,
        offset_params=offset_params,
        parent_class=parent_class,
        session=session,
        debug_num_stacks_to_drop=debug_num_stacks_to_drop,
        debug_api=debug_api,
        debug_loop=debug_loop,
    )
    if not res.is_success:
        raise JupyterAPI_Error(
            res = res
        )

    return res

# %% ../../nbs/routes/jupyter.ipynb 9
@gd.route_function
async def get_jupyter_workspace_by_id(
    auth,
    workspace_id,
    parent_class: str = None,
    session: httpx.AsyncClient = None,
    debug_api: bool = False,
    debug_num_stacks_to_drop=2,
):
    url = f"https://{auth.domo_instance}.domo.com/api/datascience/v1/workspaces/{workspace_id}"

    res = await gd.get_data(
        url=url,
        method="GET",
        auth=auth,
        parent_class=parent_class,
        session=session,
        num_stacks_to_drop=debug_num_stacks_to_drop,
        debug_api=debug_api,
    )

    if not res.is_success:
        raise JupyterAPI_Error(
            res = res
        )

    return res

# %% ../../nbs/routes/jupyter.ipynb 11
def parse_instance_service_location_and_prefix(instance: dict, domo_instance):
    url = instance["url"]

    query = urllib.parse.unquote(urllib.parse.urlparse(url).query)
    query = urllib.parse.urlparse(query.split("&")[1].replace("next=", ""))

    return {
        "service_location": query.netloc.replace(domo_instance, "")[1:],
        "service_prefix": query.path,
    }


async def get_workspace_auth_token_params(workspace_id, auth, return_raw: bool = False):
    """
    params are needed for authenticating requests inside the workspace environment
    Note: you'll also need a internally generated jupyter_token to authenticate requests
    returns { service_location , service_prefix}
    """
    res = await get_jupyter_workspace_by_id(workspace_id=workspace_id, auth=auth)

    open_instances = res.response.get("instances")

    if return_raw:
        return open_instances

    if not open_instances:
        raise JupyterAPI_WorkspaceStarted(
            res = res, message = "There are no open instances. Do you need to start the workspace?"
        )

    return parse_instance_service_location_and_prefix(
        open_instances[0], auth.domo_instance
    )

# %% ../../nbs/routes/jupyter.ipynb 13
@gd.route_function
async def start_jupyter_workspace(
    workspace_id,
    auth: dmda.DomoAuth,
    parent_class: str = None,
    session: httpx.AsyncClient = None,
    debug_num_stacks_to_drop=1,
    debug_api: bool = False,
    return_raw: bool = False,
):
    url = f"https://{auth.domo_instance}.domo.com/api/datascience/v1/workspaces/{workspace_id}/instances"

    try:
        res = await gd.get_data(
            url=url,
            method="POST",
            auth=auth,
            parent_class=parent_class,
            session=session,
            num_stacks_to_drop=debug_num_stacks_to_drop,
            debug_api=debug_api,
        )

        if return_raw:
            return res

    except RuntimeError as e:
        return rgd.ResponseGetData(
            status=500,
            response=f"starting workspace, please wait - {e}",
            is_success=False,
        )

    if res.status == 500 or res.status == 403:
        raise JupyterAPI_Error(
            res=res,
            message=f"you may not have access to this workspace {workspace_id}, is it shared with you? || OR may already be started -- {res.response}",
        )

    if not res.is_success:
        raise JupyterAPI_Error(res=res)

    res.response = "workspace started"
    return res

# %% ../../nbs/routes/jupyter.ipynb 15
@gd.route_function
async def get_jupyter_content(
    auth: dmda.DomoJupyterAuth,
    content_path: str = "",
    debug_api: bool = False,
    debug_num_stacks_to_drop=1,
    parent_class: str = None,
    session: httpx.AsyncClient = None,
):

    dmda.test_is_jupyter_auth(auth)
    
    url = f"https://{auth.domo_instance}.{auth.service_location}{auth.service_prefix}api/contents/{content_path}"

    res = await gd.get_data(
        url=f"{url}",
        method="GET",
        auth=auth,
        headers={"authorization": f"Token {auth.jupyter_token}"},
        debug_api=debug_api,
        num_stacks_to_drop=debug_num_stacks_to_drop,
        parent_class=parent_class,
        session=session,
    )
    
    if res.status == 403:
        raise JupyterAPI_Error(
            res = res, message = "unable to query API, valid jupyter_token?"
        )
    
    if not res.is_success:
        raise JupyterAPI_Error(
            res = res)

    return res

# %% ../../nbs/routes/jupyter.ipynb 20
def generate_update_jupyter_body__new_content_path(content_path):
    if not content_path:
        return ""
    
    ## replaces ./ if passed as part of url description
    if content_path.startswith('./'):
        content_path = content_path[2:]

    if "/" in content_path:
        return "/".join(content_path.split("/")[:-1])
    else:
        return ""

# %% ../../nbs/routes/jupyter.ipynb 22
def generate_update_jupyter_body__text(body, content_path=None):
    body.update(
        {
            "format": "text",
            "type": "file",
            "path": generate_update_jupyter_body__new_content_path(content_path),
        }
    )
    return body


def generate_update_jupyter_body__ipynb(body, content_path=None):

    body.update(
        {
            "format": "json",
            "type": "notebook",
            "path": generate_update_jupyter_body__new_content_path(content_path),
        }
    )
    return body


def generate_update_jupyter_body__directory(content_path, body):
    
    body.update(
        {
            "path": generate_update_jupyter_body__new_content_path(content_path),
            "format": None,
            "type": "directory",
        }
    )
    return body

# %% ../../nbs/routes/jupyter.ipynb 23
class generate_update_jupyter_body_factory(DomoEnum):
    IPYNB = partial(generate_update_jupyter_body__ipynb)
    DIRECTORY = partial(generate_update_jupyter_body__directory)
    TEXT = partial(generate_update_jupyter_body__text)
    default = partial(generate_update_jupyter_body__text)


def generate_update_jupyter_body(
    new_content, content_path: str  # my_folder/datatypes.ipynb
):
    """factory to construct properly formed body"""

    if content_path.startswith('./'):
        content_path = content_path[2:]

    content_name = os.path.normpath(content_path).split(os.sep)[-1]

    if "." in content_path:
        content_type = content_path.split(".")[-1]
    else:
        content_type = "directory"

    body = {
        "name": content_name,
        "content": new_content,
        "path": content_path,
    }
    return generate_update_jupyter_body_factory.get(content_type).value(
        body=body, content_path=content_path
    )

# %% ../../nbs/routes/jupyter.ipynb 27
@gd.route_function
async def create_jupyter_obj(
    auth: dmda.DomoJupyterAuth,
    new_content="",
    content_path: str = "",  # file name and location in jupyter
    debug_api: bool = False,
    parent_class: str = None,
    debug_num_stacks_to_drop=1,
    session: httpx.AsyncClient = None,
):
    dmda.test_is_jupyter_auth(auth)

    # removes ./ jic
    if content_path.startswith("./"):
        content_path = content_path[2:]

    body = generate_update_jupyter_body(
        new_content=new_content, content_path=content_path
    )

    content_path_split = os.path.normpath(content_path).split(os.sep)

    # new content gets created as "untitled folder" // removes the 'future name' and saves for later
    path_to_rename = content_path_split.pop(-1)

    base_url = f"https://{auth.domo_instance}.{auth.service_location}{auth.service_prefix}api/contents/"

    res_post = await gd.get_data(
        url=f"{base_url}{'/'.join(content_path_split)}",
        method="POST",
        auth=auth,
        body=body,
        debug_api=debug_api,
        parent_class=parent_class,
        num_stacks_to_drop=debug_num_stacks_to_drop,
        session=session,
    )

    if res_post.status == 403:
        raise JupyterAPI_Error(
            res = res_post, message = "unable to query API, valid jupyter_token?"
        )

    if not res_post.is_success:
        raise JupyterAPI_Error(
            res = res_post
        )

    # untitled_folder
    url = urllib.parse.urljoin(base_url, res_post.response["path"])

    # created a folder "untitled folder"
    await asyncio.sleep(3)

    res = await gd.get_data(
        url=urllib.parse.quote(url, safe="/:?=&"),
        method="PATCH",
        auth=auth,
        body={"path": content_path, 'content': new_content},
        debug_api=debug_api,
        parent_class=parent_class,
        num_stacks_to_drop=debug_num_stacks_to_drop,
        session=session,
    )

    if res.status == 403:
        raise JupyterAPI_Error(
            res = res, message = "unable to query API, valid jupyter_token?"
        )

    if res.status == 409:
        raise JupyterAPI_Error(res = res , message = "conlict during PATCH - does the content already exist?")

    if not res.is_success:
        raise JupyterAPI_Error(res=res)

    res.response = {**res_post.response, **res.response}

    return res

# %% ../../nbs/routes/jupyter.ipynb 29
@gd.route_function
async def delete_jupyter_content(
    auth: dmda.DomoJupyterAuth,
    content_path: str,  # file name and location in jupyter
    debug_api: bool = False,
    parent_class: str = None,
    debug_num_stacks_to_drop=1,
    session: httpx.AsyncClient = None,
):
    dmda.test_is_jupyter_auth(auth)

    base_url = f"https://{auth.domo_instance}.{auth.service_location}{auth.service_prefix}api/contents/"

    url = urllib.parse.urljoin(base_url, content_path)
    url=urllib.parse.quote(url, safe="/:?=&")

    res = await gd.get_data(
        url=url,
        method="DELETE",
        auth=auth,
        debug_api=debug_api,
        parent_class=parent_class,
        num_stacks_to_drop=debug_num_stacks_to_drop,
        session=session,
    )

    if res.status == 403:
        raise JupyterAPI_Error(
            res = res, message = "unable to query API, valid jupyter_token?"
        )

    if not res.is_success:
        raise JupyterAPI_Error(
            res = res
        )

    return res

# %% ../../nbs/routes/jupyter.ipynb 30
@gd.route_function
async def update_jupyter_file(
    auth: dmda.DomoJupyterAuth,
    new_content,
    content_path: str = "",  # file name and location in jupyter
    debug_api: bool = False,
    parent_class: str = None,
    debug_num_stacks_to_drop=1,
    session: httpx.AsyncClient = None,
    body : dict = None
):
    dmda.test_is_jupyter_auth(auth)

    body = body or generate_update_jupyter_body(new_content, content_path)

    content_path_split = os.path.normpath(content_path).split(os.sep)

    base_url = f"https://{auth.domo_instance}.{auth.service_location}{auth.service_prefix}api/contents/"

    url = urllib.parse.urljoin(base_url, content_path)
    url=urllib.parse.quote(url, safe="/:?=&")

    res = await gd.get_data(
        url=url,
        method="PUT",
        auth=auth,
        body=body,
        debug_api=debug_api,
        parent_class=parent_class,
        num_stacks_to_drop=debug_num_stacks_to_drop,
        session=session,
    )

    if res.status == 403:
        raise JupyterAPI_Error(
            res = res, message = "unable to query API, valid jupyter_token?"
        )

    if not res.is_success:
        raise JupyterAPI_Error(
            res = res
        )

    return res

# %% ../../nbs/routes/jupyter.ipynb 34
async def get_content_recursive(
    auth : dmda.DomoJupyterAuth,
    all_rows,
    content_path,
    logs,
    res: rgd.ResponseGetData,
    obj: dict = None,
    is_recursive: bool = True,
    is_skip_recent_executions: bool = True,
    is_skip_default_files: bool = True,
    return_raw: bool = False,
    debug_api: bool = False,
    debug_num_stacks_to_drop=0,
    parent_class=None,
    session: httpx.AsyncClient = None,
):
    # set path (on initial execution there is no object)

    dmda.test_is_jupyter_auth(auth)

    res = await get_jupyter_content(
        auth=auth,
        content_path=content_path,
        debug_api=debug_api,
        parent_class=parent_class,
        debug_num_stacks_to_drop=debug_num_stacks_to_drop + 1,
        session=session,
    )
    if return_raw:
        return res

    obj = res.response
    obj_name = obj["name"]
    obj_type = obj["type"]
    obj_path = obj["path"]
    obj_content = obj["content"] or []

    s = {"content_path": obj_path, "type": obj_type, "name": obj_name}

    if (is_skip_recent_executions and obj_path.startswith("recent_executions")) or (
        is_skip_default_files and obj_path.startswith("domo_jupyter_examples")
    ):
        res.response = all_rows
        res.logs = logs
        return res

    all_rows.append(obj)
    s.update({"is_append": True})
    logs.append(s)

    res.response = all_rows
    res.logs = logs

    if obj["type"] != "directory":
        return res

    s.update({"content": len(obj_content), "all_rows": len(all_rows)})
    logs.append(s)

    res.response = all_rows
    res.logs = logs

    if not is_recursive:
        return res

    if len(obj_content) > 0:
        await dmce.gather_with_concurrency(
            *[
                get_content_recursive(
                    auth=auth,
                    content_path=content["path"],
                    all_rows=all_rows,
                    logs=logs,
                    res=res,
                    is_skip_recent_executions=is_skip_recent_executions,
                    is_skip_default_files=is_skip_default_files,
                    debug_api=debug_api,
                    debug_num_stacks_to_drop=debug_num_stacks_to_drop + 1,
                    parent_class=parent_class,
                    session=session,
                )
                for content in obj_content
            ],
            n=5
        )

    return res

# %% ../../nbs/routes/jupyter.ipynb 35
@gd.route_function
async def get_content(
    auth: dmda.DomoJupyterAuth,
    content_path="",
    is_recursive: bool = True,
    is_skip_recent_executions: bool = True,
    is_skip_default_files: bool = True,
    return_raw: bool = False,
    debug_api: bool = False,
    debug_num_stacks_to_drop=2,
    parent_class: str = None,
    session: httpx.AsyncClient = None,
):
    dmda.test_is_jupyter_auth(auth)

    all_rows = []
    logs = []
    res = None

    return await get_content_recursive(
        auth=auth,
        content_path=content_path,
        all_rows=all_rows,
        logs=logs,
        res=res,
        is_recursive=is_recursive,
        is_skip_recent_executions=is_skip_recent_executions,
        is_skip_default_files=is_skip_default_files,
        return_raw=return_raw,
        debug_api=debug_api,
        debug_num_stacks_to_drop=debug_num_stacks_to_drop,
        parent_class=parent_class,
        session=session,
    )

# %% ../../nbs/routes/jupyter.ipynb 38
@gd.route_function
async def update_jupyter_workspace_config(
    auth : dmda.DomoAuth,
    workspace_id,
    config: dict,
    parent_class: str = None,
    session: httpx.AsyncClient = None,
    debug_api: bool = False,
    debug_num_stacks_to_drop=2,
):

    url = f"https://{auth.domo_instance}.domo.com/api/datascience/v1/workspaces/{workspace_id}"

    res = await gd.get_data(
        url=url,
        method="PUT",
        auth=auth,
        body=config,
        debug_api=debug_api,
        num_stacks_to_drop=debug_num_stacks_to_drop,
        parent_class=parent_class,
        session=session,
    )

    if not res.is_success:
        raise JupyterAPI_Error(res=res, message=f"error updating {workspace_id}")

    return res
