"""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, create, batch).
Usa jupyter_nbmodel_client para comunicarse con el servidor Jupyter via WebSocket.
Para crear notebooks usa la API REST de Jupyter (urllib PUT /api/contents).
"""

import asyncio
import json
import argparse
import sys
from urllib.error import URLError, HTTPError
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
# ---------------------------------------------------------------------------


def _notebook_exists(notebook_path: str, server_url: str, token: str) -> bool:
    """Comprueba si un notebook existe en el servidor Jupyter via GET /api/contents.

    Usa GET con ?content=0 (metadata only). HEAD no es soportado universalmente
    (Jupyter 4.x devuelve 405 Method Not Allowed en /api/contents/{path}).
    """
    headers = {"Accept": "application/json"}
    if token:
        headers["Authorization"] = f"token {token}"
    check_url = f"{server_url}/api/contents/{notebook_path}?content=0"
    req = Request(check_url, headers=headers, method="GET")
    try:
        with urlopen(req, timeout=5):
            return True
    except HTTPError as e:
        if e.code == 404:
            return False
        raise


def _auto_create_notebook(notebook_path: str, server_url: str, token: str) -> bool:
    """Crea el notebook si no existe. Retorna True si fue creado."""
    if not _notebook_exists(notebook_path, server_url, token):
        jupyter_create_notebook(notebook_path, server_url=server_url, token=token)
        return True
    return False


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


async def _append_cell(
    notebook_path: str,
    source: str,
    cell_type: str,
    server_url: str,
    token: str,
) -> dict:
    _auto_create_notebook(notebook_path, server_url, token)
    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,
    }


async def _batch_write(
    notebook_path: str,
    cells: list,
    server_url: str,
    token: str,
) -> dict:
    """Anade multiples celdas en una sola conexion WebSocket."""
    _auto_create_notebook(notebook_path, server_url, token)
    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:
        for cell in cells:
            cell_type = cell.get("type", "code")
            source = cell.get("source", "")
            if cell_type == "markdown":
                nb.add_markdown_cell(source)
            else:
                nb.add_code_cell(source)
        await asyncio.sleep(2)
    return {
        "action": "batch",
        "cells_added": len(cells),
        "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)
    )


def jupyter_create_notebook(
    notebook_path: str,
    kernel_name: str = "python3",
    server_url: str = "http://localhost:8888",
    token: str = "",
    force: bool = False,
) -> dict:
    """Crea un notebook vacio (nbformat 4) via la API REST de Jupyter.

    Usa PUT /api/contents/{path} con type: notebook. Si el notebook ya existe
    y force=False, lanza un error en vez de sobreescribirlo.

    Args:
        notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
        kernel_name: Nombre del kernel a asociar (default: python3).
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion del servidor Jupyter.
        force: Si True, sobreescribe el notebook si ya existe.

    Returns:
        dict con action, notebook y created.
    """
    headers = {
        "Content-Type": "application/json",
        "Accept": "application/json",
    }
    if token:
        headers["Authorization"] = f"token {token}"

    # Verificar si ya existe (GET con content=0; HEAD no esta soportado en Jupyter 4.x)
    check_url = f"{server_url}/api/contents/{notebook_path}"
    check_req = Request(f"{check_url}?content=0", headers=headers, method="GET")
    already_exists = False
    try:
        with urlopen(check_req, timeout=5):
            already_exists = True
    except HTTPError as e:
        if e.code != 404:
            raise
    except URLError:
        raise

    if already_exists and not force:
        raise FileExistsError(
            f"Notebook ya existe: {notebook_path}. Usa force=True para sobreescribir."
        )

    kernel_display = {
        "python3": "Python 3 (ipykernel)",
        "python": "Python 3",
    }.get(kernel_name, kernel_name)

    notebook_content = {
        "nbformat": 4,
        "nbformat_minor": 5,
        "metadata": {
            "kernelspec": {
                "name": kernel_name,
                "display_name": kernel_display,
                "language": "python",
            },
            "language_info": {
                "name": "python",
            },
        },
        "cells": [],
    }

    body = json.dumps({
        "type": "notebook",
        "content": notebook_content,
    }).encode("utf-8")

    put_req = Request(check_url, data=body, headers=headers, method="PUT")
    with urlopen(put_req, timeout=10) as resp:
        resp.read()

    return {
        "action": "create",
        "notebook": notebook_path,
        "created": True,
    }


def jupyter_batch_write(
    notebook_path: str,
    cells: list,
    server_url: str = "http://localhost:8888",
    token: str = "",
) -> dict:
    """Anade multiples celdas al notebook en una sola conexion WebSocket.

    Mucho mas eficiente que llamar append-code/append-markdown N veces porque
    abre una unica conexion WebSocket y hace un solo sleep de sincronizacion
    al final de todas las inserciones.

    Args:
        notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
        cells: Lista de dicts con claves "type" ("code"|"markdown") y "source" (str).
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion del servidor Jupyter.

    Returns:
        dict con action, cells_added y notebook.
    """
    return asyncio.run(_batch_write(notebook_path, cells, 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)

    # create
    p_cr = sub.add_parser("create", help="Crea un notebook vacio via API REST")
    p_cr.add_argument("notebook", help="Ruta del notebook a crear")
    p_cr.add_argument("--kernel", default="python3", help="Nombre del kernel (default: python3)")
    p_cr.add_argument("--force", action="store_true", help="Sobreescribir si ya existe")
    add_common(p_cr)

    # batch
    p_bt = sub.add_parser("batch", help="Anade multiples celdas en una sola conexion WebSocket")
    p_bt.add_argument("notebook", help="Ruta del notebook")
    p_bt.add_argument(
        "--from",
        dest="cells_source",
        default="-",
        metavar="FILE",
        help="Archivo JSON con las celdas, o '-' para leer de stdin (default: -)",
    )
    add_common(p_bt)

    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)
    elif args.command == "create":
        result = jupyter_create_notebook(
            args.notebook,
            kernel_name=args.kernel,
            server_url=args.server,
            token=args.token,
            force=args.force,
        )
    elif args.command == "batch":
        if args.cells_source == "-":
            raw = sys.stdin.read()
        else:
            with open(args.cells_source, "r", encoding="utf-8") as f:
                raw = f.read()
        cells = json.loads(raw)
        result = jupyter_batch_write(args.notebook, cells, args.server, args.token)
    else:
        parser.print_help()
        return

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


if __name__ == "__main__":
    main()