Update src/services/chart_browser_engine.py

src/services/chart_browser_engine.py

@@ -1,3 +1,4 @@
+import io
 import json
 import os
 import re
@@ -11,15 +12,16 @@ from google.auth.transport.requests import Request as GoogleAuthRequest
 from google.oauth2.credentials import Credentials
 from googleapiclient.discovery import build
 from googleapiclient.errors import HttpError
-from googleapiclient.http import MediaIoBaseDownload
+from googleapiclient.http import MediaIoBaseDownload, MediaIoBaseUpload
 from PIL import Image, ImageOps
 
 from ..config import update_user_keys, get_user_keys, firestore
 from ..state import get_user_temp_dir
 from .journal_engine import FOLDER_MIME
 
-BROWSER_SCOPES = ['https://www.googleapis.com/auth/drive.readonly']
+READONLY_SCOPE = 'https://www.googleapis.com/auth/drive.readonly'
 FULL_DRIVE_SCOPE = 'https://www.googleapis.com/auth/drive'
+BROWSER_SCOPES = [FULL_DRIVE_SCOPE]
 
 TOKEN_FIELD = "google_browser_token_json"
 LAST_FOLDER_FIELD = "chart_browser_last_folder"
@@ -40,18 +42,39 @@ ID_RE = re.compile(r"^[A-Za-z0-9_-]{1,128}$")
 
 MAX_IMAGES = 1000
 MAX_FILE_BYTES = 30 * 1024 * 1024
+MAX_UPLOAD_BYTES = 20 * 1024 * 1024
+MAX_NAME_CHARS = 120
+RESUMABLE_ABOVE = 5 * 1024 * 1024
+
+# PIL format name -> stored mime.
+PIL_FORMAT_MIME = {
+    "JPEG": "image/jpeg",
+    "PNG": "image/png",
+    "WEBP": "image/webp",
+    "GIF": "image/gif",
+    "BMP": "image/bmp",
+    "AVIF": "image/avif",
+}
+MIME_EXTS = {
+    "image/jpeg": (".jpg", ".jpeg"),
+    "image/png": (".png",),
+    "image/webp": (".webp",),
+    "image/gif": (".gif",),
+    "image/bmp": (".bmp",),
+    "image/avif": (".avif",),
+}
+IMAGE_FIELDS = ("id,name,mimeType,size,modifiedTime,webViewLink,"
+                "imageMediaMetadata(width,height)")
 MAX_CACHE_BYTES = 300 * 1024 * 1024
 CACHE_TTL_SECONDS = 48 * 60 * 60
 DOWNLOAD_CHUNK = 4 * 1024 * 1024
 
 PREVIEW_MAX_PX = 1600
 PREVIEW_QUALITY = 78
-PREVIEW_MIN_BYTES = 250 * 1024   # smaller originals are served as-is
-PREVIEW_TAG = "~p"               # '~' can never appear in a Drive id
+PREVIEW_MIN_BYTES = 250 * 1024  
+PREVIEW_TAG = "~p"               
 
-# Drive's orderBy for each sort. Both date sorts fetch newest-first so that,
-# in a folder over MAX_IMAGES, the cut-off drops the oldest files and never
-# the latest ones. The final order is applied in list_images().
+# Drive's orderBy for each sort. 
 SORT_ORDERS = {
     "name": "name_natural",
     "newest": "modifiedTime desc,name_natural desc",
@@ -85,6 +108,10 @@ class TooLarge(Exception):
     pass
 
 
+class InvalidUpload(Exception):
+    pass
+
+
 def valid_id(value) -> bool:
     return isinstance(value, str) and bool(ID_RE.match(value))
 
@@ -107,7 +134,20 @@ class ChartBrowserEngine:
         if not creds:
             return False
         scopes = set(getattr(creds, 'scopes', None) or [])
-        return bool(scopes & {BROWSER_SCOPES[0], FULL_DRIVE_SCOPE})
+        return bool(scopes & {READONLY_SCOPE, FULL_DRIVE_SCOPE})
+
+    @staticmethod
+    def has_write_scope(creds) -> bool:
+        return FULL_DRIVE_SCOPE in set(getattr(creds, 'scopes', None) or [])
+
+    @staticmethod
+    def token_can_write(token_json) -> bool:
+        """Cheap check on the stored token (no refresh, no network) for the page."""
+        try:
+            scopes = json.loads(token_json).get('scopes') or []
+        except Exception:
+            return False
+        return FULL_DRIVE_SCOPE in scopes
 
     @staticmethod
     def get_creds(uid, user_data=None):
@@ -248,37 +288,172 @@ class ChartBrowserEngine:
         return folders
 
     @staticmethod
-    def list_images(service, folder_id, sort=DEFAULT_SORT):
-        mime_clause = " or ".join(f"mimeType='{m}'" for m in IMAGE_MIME_EXT)
-        q = f"'{folder_id}' in parents and trashed=false and ({mime_clause})"
-        fields = ("id,name,mimeType,size,modifiedTime,webViewLink,"
-                  "imageMediaMetadata(width,height)")
-        raw, truncated = ChartBrowserEngine._list_all(
-            service, q, fields, SORT_ORDERS.get(sort, SORT_ORDERS[DEFAULT_SORT]), LIST_PAGE_SIZE, MAX_IMAGES
-        )
-        images = []
-        for f in raw:
-            meta = f.get('imageMediaMetadata') or {}
-            try:
-                size = int(f.get('size') or 0)
-            except (TypeError, ValueError):
-                size = 0
-            images.append({
-                "id": f['id'],
-                "name": f.get('name', f['id']),
-                "mime": f.get('mimeType'),
-                "size": size,
-                "modified": f.get('modifiedTime'),
-                "width": meta.get('width'),
-                "height": meta.get('height'),
-                "link": f.get('webViewLink'),
-            })
+    def _entry(f) -> dict:
+        meta = f.get('imageMediaMetadata') or {}
+        try:
+            size = int(f.get('size') or 0)
+        except (TypeError, ValueError):
+            size = 0
+        return {
+            "id": f['id'],
+            "name": f.get('name', f['id']),
+            "mime": f.get('mimeType'),
+            "size": size,
+            "modified": f.get('modifiedTime'),
+            "width": meta.get('width'),
+            "height": meta.get('height'),
+            "link": f.get('webViewLink'),
+        }
+
+    @staticmethod
+    def _order(images, sort):
         if sort in ("oldest", "newest"):
             # Sort here as well so the order never depends on Drive's tie
             # handling: by modified time, then name for identical timestamps.
             images.sort(key=lambda i: (i.get("modified") or "", i["name"].lower()),
                         reverse=(sort == "newest"))
-        return images, truncated
+        return images
+
+    @staticmethod
+    def list_images(service, folder_id, sort=DEFAULT_SORT):
+        mime_clause = " or ".join(f"mimeType='{m}'" for m in IMAGE_MIME_EXT)
+        q = f"'{folder_id}' in parents and trashed=false and ({mime_clause})"
+        raw, truncated = ChartBrowserEngine._list_all(
+            service, q, IMAGE_FIELDS, SORT_ORDERS.get(sort, SORT_ORDERS[DEFAULT_SORT]), LIST_PAGE_SIZE, MAX_IMAGES
+        )
+        images = [ChartBrowserEngine._entry(f) for f in raw]
+        return ChartBrowserEngine._order(images, sort), truncated
+
+    # --- Upload / delete ---
+    @staticmethod
+    def _clean_name(raw, mime) -> str:
+        name = (raw or "").replace("\\", "/").split("/")[-1]
+        name = re.sub(r"[\x00-\x1f\x7f]", "", name).strip().strip(".")
+        stem, dot, ext = name.rpartition(".")
+        if not dot:
+            stem, ext = name, ""
+        allowed = MIME_EXTS[mime]
+        if ("." + ext.lower()) not in allowed:
+            stem, ext = name, allowed[0].lstrip(".")
+        stem = stem.strip() or "chart"
+        keep = max(1, MAX_NAME_CHARS - len(ext) - 1)
+        return f"{stem[:keep]}.{ext}"
+
+    @staticmethod
+    def _read_image(stream):
+        """Read the upload (bounded) and identify it from its bytes."""
+        data = stream.read(MAX_UPLOAD_BYTES + 1)
+        if not data:
+            raise InvalidUpload("The file is empty.")
+        if len(data) > MAX_UPLOAD_BYTES:
+            raise InvalidUpload(f"File is too large (max {MAX_UPLOAD_BYTES // (1024 * 1024)} MB).")
+        try:
+            with Image.open(io.BytesIO(data)) as im:
+                fmt = im.format
+                im.verify()
+        except Exception:
+            raise InvalidUpload("That file is not a valid image.")
+        mime = PIL_FORMAT_MIME.get(fmt)
+        if not mime:
+            raise InvalidUpload("Unsupported image type. Use PNG, JPG, WebP, GIF, BMP or AVIF.")
+        return data, mime
+
+    @staticmethod
+    def _refresh_listing(uid, service, folder_id, sort, add=None, remove_id=None):
+        """Re-list Drive after a change and save the index. add/remove_id patch
+        the result in case Drive's listing lags behind our own write."""
+        idx = ChartBrowserEngine.load_index(uid)
+        if (idx.get("folder") or {}).get("id") == folder_id and "parent" in idx:
+            folder = {"id": folder_id, "name": idx["folder"]["name"], "parent": idx["parent"]}
+        else:
+            folder = ChartBrowserEngine.folder_info(service, folder_id, uid)
+
+        images, truncated = ChartBrowserEngine.list_images(service, folder_id, sort)
+        if remove_id:
+            images = [i for i in images if i["id"] != remove_id]
+        if add and not truncated and all(i["id"] != add["id"] for i in images):
+            images.append(add)
+            ChartBrowserEngine._order(images, sort)
+        ChartBrowserEngine.save_index(
+            uid, {"id": folder["id"], "name": folder["name"]}, images,
+            sort=sort, parent=folder["parent"], truncated=truncated)
+        return folder, images, truncated
+
+    @staticmethod
+    def upload_image(uid, service, folder_id, stream, filename, sort=DEFAULT_SORT):
+        if folder_id != 'root' and not valid_id(folder_id):
+            raise NotFound("Bad folder")
+        data, mime = ChartBrowserEngine._read_image(stream)
+        name = ChartBrowserEngine._clean_name(filename, mime)
+
+        media = MediaIoBaseUpload(io.BytesIO(data), mimetype=mime,
+                                  resumable=len(data) > RESUMABLE_ABOVE)
+        try:
+            created = service.files().create(
+                body={'name': name, 'parents': [folder_id]},
+                media_body=media,
+                fields=IMAGE_FIELDS,
+                supportsAllDrives=True,
+            ).execute()
+        except HttpError as e:
+            if e.resp.status == 404:
+                raise NotFound("Folder not found")
+            raise
+        entry = ChartBrowserEngine._entry(created)
+
+        folder, images, truncated = ChartBrowserEngine._refresh_listing(
+            uid, service, folder_id, sort, add=entry)
+
+        d = ChartBrowserEngine.cache_dir(uid)
+        final = d / f"{entry['id']}{IMAGE_MIME_EXT[mime]}"
+        part = d / f"{entry['id']}.part"
+        try:
+            with open(part, "wb") as fh:
+                fh.write(data)
+            os.replace(part, final)
+            ChartBrowserEngine._enforce_cap(d, final.name)
+        except OSError:
+            pass
+        finally:
+            if part.exists():
+                try:
+                    part.unlink()
+                except OSError:
+                    pass
+        return entry, folder, images, truncated
+
+    @staticmethod
+    def delete_image(uid, service, file_id):
+        """Move an image to the Drive trash. Only files in the folder that is
+        currently open (the saved index) can be deleted."""
+        if not valid_id(file_id):
+            raise NotFound("Bad id")
+        idx = ChartBrowserEngine.load_index(uid)
+        if not any(i["id"] == file_id for i in idx.get("images", [])):
+            raise NotFound("Image is not in the open folder")
+        folder_id = idx["folder"]["id"]
+        sort = idx.get("sort") or DEFAULT_SORT
+
+        try:
+            service.files().update(
+                fileId=file_id, body={'trashed': True},
+                fields='id', supportsAllDrives=True,
+            ).execute()
+        except HttpError as e:
+            if e.resp.status != 404:   # already gone counts as deleted
+                raise
+
+        try:
+            return ChartBrowserEngine._refresh_listing(
+                uid, service, folder_id, sort, remove_id=file_id)
+        except Exception as e:
+            print(f"Chart Browser relist after delete failed: {e}")
+            images = [i for i in idx["images"] if i["id"] != file_id]
+            folder = {"id": folder_id, "name": idx["folder"]["name"], "parent": idx.get("parent")}
+            ChartBrowserEngine.save_index(
+                uid, {"id": folder_id, "name": folder["name"]}, images,
+                sort=sort, parent=folder["parent"], truncated=bool(idx.get("truncated")))
+            return folder, images, bool(idx.get("truncated"))
 
     # --- Temp cache ---
     @staticmethod