import os
import httpx
import nbformat
from nbformat import NotebookNode, ValidationError
from nbconvert import HTMLExporter
from starlette.applications import Starlette
from starlette.exceptions import HTTPException
from starlette.responses import FileResponse, JSONResponse, HTMLResponse
from starlette.requests import Request
from starlette.routing import Route

client = httpx.AsyncClient()
html_exporter = HTMLExporter(template_name="lab")


async def homepage(_):
    return FileResponse("static/index.html")


async def healthz(_):
    return JSONResponse({"success": True})


async def convert(req: Request):
    url = req.query_params.get("url")
    if not url:
        raise HTTPException(400, "Param url is missing")
    print("\n===", url)
    r = await client.get(
        url,
        headers={"Authorization": f"Bearer {os.environ.get('HF_TOKEN')}"},
        follow_redirects=True,
        # httpx no follow redirect by default
    )
    if r.status_code != 200:
        raise HTTPException(
            400, f"Got an error {r.status_code} when fetching remote file"
        )

    # Capture potential validation error:
    try:
        notebook_node: NotebookNode = nbformat.reads(
            r.text,
            as_version=nbformat.current_nbformat,
        )
    except nbformat.reader.NotJSONError:
        print(400, f"Notebook is not JSON ({url})")
        raise HTTPException(400, f"Notebook is not JSON ({url})")
    except ValidationError as e:
        print(
            400,
            f"Notebook is invalid according to nbformat: {e} ({url})",
        )
        raise HTTPException(
            400,
            f"Notebook is invalid according to nbformat: {e} ({url})",
        )

    print(f"Input: nbformat v{notebook_node.nbformat}.{notebook_node.nbformat_minor}")
    body, _ = html_exporter.from_notebook_node(notebook_node)
    # TODO(customize or simplify template?)
    # TODO(also check source code for jupyter/nbviewer)
    return HTMLResponse(body)


app = Starlette(
    debug=False,
    routes=[
        Route("/", homepage),
        Route("/healthz", healthz),
        Route("/convert", convert),
    ],
)