From 9a1300f565ec9ff9c801d50bbabe1bfc3c0b99bf Mon Sep 17 00:00:00 2001
From: Kaustubh Maske Patil <37668193+nikochiko@users.noreply.github.com>
Date: Thu, 19 Oct 2023 04:16:26 +0530
Subject: [PATCH] Add better error handling for UI

---
 routers/root.py               | 19 +++++++++++++++++++
 templates/404.html            | 11 -----------
 templates/errors/404.html     |  8 ++++++++
 templates/errors/unknown.html |  8 ++++++++
 4 files changed, 35 insertions(+), 11 deletions(-)
 delete mode 100644 templates/404.html
 create mode 100644 templates/errors/404.html
 create mode 100644 templates/errors/unknown.html

diff --git a/routers/root.py b/routers/root.py
index 42627572e..2acd63a5b 100644
--- a/routers/root.py
+++ b/routers/root.py
@@ -72,6 +72,25 @@ async def favicon():
     return FileResponse("static/favicon.ico")
 
 
+@app.post("/handleError/")
+async def handle_error(request: Request):
+    context = {"request": request, "settings": settings}
+
+    def not_found():
+        st.html(templates.get_template("errors/404.html").render(**context))
+
+    def unknown_error():
+        st.html(templates.get_template("errors/unknown.html").render(**context))
+
+    body = await request.json()
+
+    match body["status"]:
+        case 404:
+            return st.runner(lambda: page_wrapper(request, not_found))
+        case _:
+            return st.runner(lambda: page_wrapper(request, unknown_error))
+
+
 @app.get("/login/")
 def login(request: Request):
     if request.user and not request.user.is_anonymous:
diff --git a/templates/404.html b/templates/404.html
deleted file mode 100644
index 7389e5072..000000000
--- a/templates/404.html
+++ /dev/null
@@ -1,11 +0,0 @@
-{% extends 'base.html' %}
-
-{% block content %}
-    {% include 'login_container.html' %}
-    <div style="text-align: center;">
-        <h1>
-            Oops, we can't find this page!
-        </h1>
-        <img src="https://media.tenor.com/IHdlTRsmcS4AAAAC/404.gif" height="404" alt="404">
-    </div>
-{% endblock content %}
\ No newline at end of file
diff --git a/templates/errors/404.html b/templates/errors/404.html
new file mode 100644
index 000000000..2b3bcc4a5
--- /dev/null
+++ b/templates/errors/404.html
@@ -0,0 +1,8 @@
+{% block content %}
+<div style="text-align: center; margin: 2rem 0;">
+    <h1>
+        Oops, we can't find this page!
+    </h1>
+    <img src="https://media.tenor.com/IHdlTRsmcS4AAAAC/404.gif" height="404" alt="404">
+</div>
+{% endblock content %}
diff --git a/templates/errors/unknown.html b/templates/errors/unknown.html
new file mode 100644
index 000000000..c66925a2e
--- /dev/null
+++ b/templates/errors/unknown.html
@@ -0,0 +1,8 @@
+{% block content %}
+<div style="text-align: center; margin: 2rem 0;">
+    <h1>
+        Uh oh... Something went wrong
+    </h1>
+    <img src="https://media.tenor.com/1Wlr9trTPLwAAAAd/robotically-robot.gif" height="404" alt="404">
+</div>
+{% endblock %}