"""Operaciones de escritura sobre celdas de un notebook Jupyter via colaboracion en tiempo real.

NO ejecuta celdas — solo modifica la estructura del notebook (append, insert, edit, delete).
Usa jupyter_nbmodel_client para comunicarse con el servidor Jupyter via WebSocket.
"""

import asyncio
import json
import argparse
from urllib.error import URLError
from urllib.request import Request, urlopen

from jupyter_nbmodel_client import NbModelClient, get_jupyter_notebook_websocket_url


def _resolve_collab_username(server_url: str, token: str) -> str:
    """Resolve the display name of the active user in Jupyter collaboration."""
    headers = {"Accept": "application/json"}
    if token:
        headers["Authorization"] = f"token {token}"
    try:
        req = Request(f"{server_url}/api/me", headers=headers)
        with urlopen(req, timeout=5) as resp:
            me = json.loads(resp.read())
        identity = me.get("identity", {})
        return identity.get("display_name", "") or identity.get("username", "") or identity.get("name", "Anonymous")
    except (URLError, OSError, json.JSONDecodeError):
        return "Anonymous"


# ---------------------------------------------------------------------------
# Helpers internos async
# ---------------------------------------------------------------------------


async def _append_cell(
    notebook_path: str,
    source: str,
    cell_type: str,
    server_url: str,
    token: str,
) -> dict:
    ws_url = get_jupyter_notebook_websocket_url(
        server_url=server_url,
        token=token,
        path=notebook_path,
    )
    username = _resolve_collab_username(server_url, token)
    async with NbModelClient(ws_url, username=username) as nb:
        if cell_type == "markdown":
            nb.add_markdown_cell(source)
        else:
            nb.add_code_cell(source)
        cell_index = len(nb) - 1
        await asyncio.sleep(2)
    return {
        "action": f"append_{cell_type}",
        "cell_index": cell_index,
        "notebook": notebook_path,
    }


async def _insert_cell(
    notebook_path: str,
    cell_index: int,
    source: str,
    cell_type: str,
    server_url: str,
    token: str,
) -> dict:
    ws_url = get_jupyter_notebook_websocket_url(
        server_url=server_url,
        token=token,
        path=notebook_path,
    )
    username = _resolve_collab_username(server_url, token)
    async with NbModelClient(ws_url, username=username) as nb:
        nb.insert_cell(cell_index, cell_type=cell_type, source=source)
        await asyncio.sleep(2)
    return {
        "action": "insert",
        "cell_index": cell_index,
        "cell_type": cell_type,
        "notebook": notebook_path,
    }


async def _edit_cell(
    notebook_path: str,
    cell_index: int,
    source: str,
    server_url: str,
    token: str,
) -> dict:
    ws_url = get_jupyter_notebook_websocket_url(
        server_url=server_url,
        token=token,
        path=notebook_path,
    )
    username = _resolve_collab_username(server_url, token)
    async with NbModelClient(ws_url, username=username) as nb:
        nb.set_cell_source(cell_index, source)
        await asyncio.sleep(2)
    return {
        "action": "edit",
        "cell_index": cell_index,
        "notebook": notebook_path,
    }


async def _delete_cell(
    notebook_path: str,
    cell_index: int,
    server_url: str,
    token: str,
) -> dict:
    ws_url = get_jupyter_notebook_websocket_url(
        server_url=server_url,
        token=token,
        path=notebook_path,
    )
    username = _resolve_collab_username(server_url, token)
    async with NbModelClient(ws_url, username=username) as nb:
        nb.delete_cell(cell_index)
        await asyncio.sleep(2)
    return {
        "action": "delete",
        "cell_index": cell_index,
        "notebook": notebook_path,
    }


# ---------------------------------------------------------------------------
# API publica sincrona
# ---------------------------------------------------------------------------


def jupyter_append_code(
    notebook_path: str,
    source: str,
    server_url: str = "http://localhost:8888",
    token: str = "",
) -> dict:
    """Anade una celda de codigo al final del notebook.

    Args:
        notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
        source: Codigo fuente de la celda.
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion del servidor Jupyter.

    Returns:
        dict con action, cell_index y notebook.
    """
    return asyncio.run(_append_cell(notebook_path, source, "code", server_url, token))


def jupyter_append_markdown(
    notebook_path: str,
    source: str,
    server_url: str = "http://localhost:8888",
    token: str = "",
) -> dict:
    """Anade una celda markdown al final del notebook.

    Args:
        notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
        source: Contenido markdown de la celda.
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion del servidor Jupyter.

    Returns:
        dict con action, cell_index y notebook.
    """
    return asyncio.run(
        _append_cell(notebook_path, source, "markdown", server_url, token)
    )


def jupyter_insert_cell(
    notebook_path: str,
    cell_index: int,
    source: str,
    cell_type: str = "code",
    server_url: str = "http://localhost:8888",
    token: str = "",
) -> dict:
    """Inserta una celda en una posicion especifica del notebook.

    Args:
        notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
        cell_index: Indice donde insertar (0 = primer posicion).
        source: Contenido de la celda.
        cell_type: Tipo de celda: "code" o "markdown".
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion del servidor Jupyter.

    Returns:
        dict con action, cell_index, cell_type y notebook.
    """
    return asyncio.run(
        _insert_cell(notebook_path, cell_index, source, cell_type, server_url, token)
    )


def jupyter_edit_cell(
    notebook_path: str,
    cell_index: int,
    source: str,
    server_url: str = "http://localhost:8888",
    token: str = "",
) -> dict:
    """Sobrescribe el contenido de una celda existente.

    Args:
        notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
        cell_index: Indice de la celda a editar (0-based).
        source: Nuevo contenido de la celda.
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion del servidor Jupyter.

    Returns:
        dict con action, cell_index y notebook.
    """
    return asyncio.run(
        _edit_cell(notebook_path, cell_index, source, server_url, token)
    )


def jupyter_delete_cell(
    notebook_path: str,
    cell_index: int,
    server_url: str = "http://localhost:8888",
    token: str = "",
) -> dict:
    """Elimina una celda del notebook.

    Args:
        notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
        cell_index: Indice de la celda a eliminar (0-based).
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion del servidor Jupyter.

    Returns:
        dict con action, cell_index y notebook.
    """
    return asyncio.run(
        _delete_cell(notebook_path, cell_index, server_url, token)
    )


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


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="jupyter_write",
        description="Operaciones de escritura sobre celdas de un notebook Jupyter.",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    # Argumentos comunes
    def add_common(p: argparse.ArgumentParser) -> None:
        p.add_argument("--server", default="http://localhost:8888", help="URL del servidor Jupyter")
        p.add_argument("--token", default="", help="Token de autenticacion")

    # append-code
    p_ac = sub.add_parser("append-code", help="Anade celda de codigo al final")
    p_ac.add_argument("notebook", help="Ruta del notebook")
    p_ac.add_argument("source", help="Codigo fuente")
    add_common(p_ac)

    # append-markdown
    p_am = sub.add_parser("append-markdown", help="Anade celda markdown al final")
    p_am.add_argument("notebook", help="Ruta del notebook")
    p_am.add_argument("source", help="Contenido markdown")
    add_common(p_am)

    # insert
    p_ins = sub.add_parser("insert", help="Inserta celda en posicion especifica")
    p_ins.add_argument("notebook", help="Ruta del notebook")
    p_ins.add_argument("index", type=int, help="Indice de insercion (0-based)")
    p_ins.add_argument("source", help="Contenido de la celda")
    p_ins.add_argument("--type", dest="cell_type", choices=["code", "markdown"], default="code")
    add_common(p_ins)

    # edit
    p_ed = sub.add_parser("edit", help="Sobrescribe el contenido de una celda")
    p_ed.add_argument("notebook", help="Ruta del notebook")
    p_ed.add_argument("index", type=int, help="Indice de la celda (0-based)")
    p_ed.add_argument("source", help="Nuevo contenido")
    add_common(p_ed)

    # delete
    p_del = sub.add_parser("delete", help="Elimina una celda")
    p_del.add_argument("notebook", help="Ruta del notebook")
    p_del.add_argument("index", type=int, help="Indice de la celda (0-based)")
    add_common(p_del)

    return parser


def main() -> None:
    parser = _build_parser()
    args = parser.parse_args()

    if args.command == "append-code":
        result = jupyter_append_code(args.notebook, args.source, args.server, args.token)
    elif args.command == "append-markdown":
        result = jupyter_append_markdown(args.notebook, args.source, args.server, args.token)
    elif args.command == "insert":
        result = jupyter_insert_cell(
            args.notebook, args.index, args.source, args.cell_type, args.server, args.token
        )
    elif args.command == "edit":
        result = jupyter_edit_cell(args.notebook, args.index, args.source, args.server, args.token)
    elif args.command == "delete":
        result = jupyter_delete_cell(args.notebook, args.index, args.server, args.token)
    else:
        parser.print_help()
        return

    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()