From 81ba2cd2d23f38f73d59d3fa21c4d59b2a0501e0 Mon Sep 17 00:00:00 2001
From: arkohut <39525455+arkohut@users.noreply.github.com>
Date: Mon, 9 Sep 2024 20:58:06 +0800
Subject: [PATCH] feat: add default library and command shortcuts for it

---
 memos/commands.py         | 72 +++++++++++++++++++++++++++++++++++++++
 memos/config.py           |  1 +
 screen_recorder/record.py |  6 ++--
 3 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/memos/commands.py b/memos/commands.py
index 00753e8..8092918 100644
--- a/memos/commands.py
+++ b/memos/commands.py
@@ -762,5 +762,77 @@ def init():
         print("Initialization failed. Please check the error messages above.")
 
 
+@app.command("scan")
+def scan_default_library():
+    """
+    Scan the screenshots directory and add it to the library if empty.
+    """
+    # Get the default library
+    response = httpx.get(f"{BASE_URL}/libraries")
+    if response.status_code != 200:
+        print(f"Failed to retrieve libraries: {response.status_code} - {response.text}")
+        return
+
+    libraries = response.json()
+    default_library = next(
+        (lib for lib in libraries if lib["name"] == settings.default_library), None
+    )
+
+    if not default_library:
+        # Create the default library if it doesn't exist
+        response = httpx.post(
+            f"{BASE_URL}/libraries",
+            json={"name": settings.default_library, "folders": []},
+        )
+        if response.status_code != 200:
+            print(
+                f"Failed to create default library: {response.status_code} - {response.text}"
+            )
+            return
+        default_library = response.json()
+
+    # Check if the library is empty
+    if not default_library["folders"]:
+        # Add the screenshots directory to the library
+        screenshots_dir = Path(settings.screenshots_dir).resolve()
+        response = httpx.post(
+            f"{BASE_URL}/libraries/{default_library['id']}/folders",
+            json={"folders": [str(screenshots_dir)]},
+        )
+        if response.status_code != 200:
+            print(
+                f"Failed to add screenshots directory: {response.status_code} - {response.text}"
+            )
+            return
+        print(f"Added screenshots directory: {screenshots_dir}")
+
+    # Scan the library
+    print(f"Scanning library: {default_library['name']}")
+    scan(default_library["id"], plugins=None, folders=None)
+
+
+@app.command("index")
+def index_default_library():
+    """
+    Index the default library for memos.
+    """
+    # Get the default library
+    response = httpx.get(f"{BASE_URL}/libraries")
+    if response.status_code != 200:
+        print(f"Failed to retrieve libraries: {response.status_code} - {response.text}")
+        return
+
+    libraries = response.json()
+    default_library = next(
+        (lib for lib in libraries if lib["name"] == settings.default_library), None
+    )
+
+    if not default_library:
+        print("Default library does not exist.")
+        return
+
+    index(default_library["id"], force=False, folders=None)
+
+
 if __name__ == "__main__":
     app()
diff --git a/memos/config.py b/memos/config.py
index db5ffbc..309a373 100644
--- a/memos/config.py
+++ b/memos/config.py
@@ -49,6 +49,7 @@ class Settings(BaseSettings):
     base_dir: str = str(Path.home() / ".memos")
     database_path: str = os.path.join(base_dir, "database.db")
     default_library: str = "screenshots"
+    screenshots_dir: str = os.path.join(base_dir, "screenshots")
     
     typesense_host: str = "localhost"
     typesense_port: str = "8108"
diff --git a/screen_recorder/record.py b/screen_recorder/record.py
index 8bc101a..ba141a0 100644
--- a/screen_recorder/record.py
+++ b/screen_recorder/record.py
@@ -9,6 +9,8 @@ from screen_recorder.common import (
     take_screenshot,
     is_screen_locked,
 )
+from pathlib import Path
+from memos.config import settings
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
@@ -51,12 +53,12 @@ def main():
         "--threshold", type=int, default=4, help="Threshold for image similarity"
     )
     parser.add_argument(
-        "--base-dir", type=str, default="~/tmp", help="Base directory for screenshots"
+        "--base-dir", type=str, help="Base directory for screenshots"
     )
     parser.add_argument("--once", action="store_true", help="Run once and exit")
     args = parser.parse_args()
 
-    base_dir = os.path.expanduser(args.base_dir)
+    base_dir = os.path.expanduser(args.base_dir) if args.base_dir else settings.screenshots_dir
     previous_hashes = load_previous_hashes(base_dir)
 
     if args.once: