From 0bc8ce39f870a8f4c351a277aa7717c4cb3d5640 Mon Sep 17 00:00:00 2001
From: Ovler <ovlertheyoung@gmail.com>
Date: Sat, 15 Jun 2024 14:09:15 +0800
Subject: [PATCH] feat: Add API endpoints for archiving from user, favlist, and
 collection

The code changes add new API endpoints for archiving videos from user, favlist, and collection sources. It includes the necessary functions to run the corresponding commands and process the output. This feature enhances the functionality of the application by allowing users to easily archive videos from different sources.
---
 biliarchiver/rest_api/main.py | 73 ++++++++++++++++++++++++++++++++++-
 1 file changed, 72 insertions(+), 1 deletion(-)

diff --git a/biliarchiver/rest_api/main.py b/biliarchiver/rest_api/main.py
index b6ccf01..86a461f 100644
--- a/biliarchiver/rest_api/main.py
+++ b/biliarchiver/rest_api/main.py
@@ -3,9 +3,11 @@ from asyncio import Queue
 from contextlib import asynccontextmanager
 from datetime import datetime
 from typing import List
+import subprocess
+import re
 
 try:
-    from fastapi import FastAPI
+    from fastapi import FastAPI, HTTPException
 except ImportError:
     print("Please install fastapi")
     print("`pip install fastapi`")
@@ -102,6 +104,75 @@ async def delete(vid: str):
 
     return {"success": False, "vid": vid, "status": "not_found"}
 
+@app.get("/archive/user/{userid}")
+@app.post("/archive/user/{userid}")
+async def add_from_user(userid: int):
+    return await add_from_source("user", userid)
+
+@app.get("/archive/favlist/{favlistid}")
+@app.post("/archive/favlist/{favlistid}")
+async def add_from_favlist(favlistid: int):
+    return await add_from_source("favlist", favlistid)
+
+@app.get("/archive/collection/{sid}")
+@app.post("/archive/collection/{sid}")
+async def add_from_collection(sid: int):
+    return await add_from_source("collection", sid)
+
+
+async def add_from_source(source_type: str, source_id: int):
+    # make sure source_id is valid integer
+    try: 
+        source_id = int(source_id)
+    except ValueError:
+        raise HTTPException(status_code=400, detail="Invalid source id")
+
+    cmd_mapping = {
+        "user": ["biliarchiver", "get", "--up-videos", source_id],
+        "favlist": ["biliarchiver", "get", "--favlist", source_id],
+        "collection": ["biliarchiver", "get", "--series", source_id],
+    }
+
+    print(f"Adding {source_type} {source_id} to queue...")
+
+    if source_type not in cmd_mapping:
+        raise HTTPException(status_code=400, detail="Invalid source type")
+
+    cmd = cmd_mapping[source_type]
+
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        print(result.stdout)
+        output = result.stdout
+        txt_path = parse_txt_path_from_output(output)
+        print(f"Txt file path: {txt_path}")
+        bvids = read_bvids_from_txt(txt_path)
+        print(f"bvids: {bvids}")
+        for bvid in bvids:
+            video = BiliVideo(bvid, status=VideoStatus.pending)
+            await pending_queue.put(video)
+        return {"success": True, "source_type": source_type, "source_id": source_id, "bvids": bvids}
+    except subprocess.CalledProcessError as e:
+        raise HTTPException(status_code=500, detail=f"Command failed: {e.stderr}")
+
+
+def parse_txt_path_from_output(output: str) -> str:
+    # Use a regular expression that can match both Windows and Linux file paths
+    matches = re.findall(r"\s(([A-Za-z]:)?[\\/].+\.txt)", output)
+    if not matches:
+        raise HTTPException(status_code=500, detail="Failed to parse txt file path from output")
+    # Select the last match
+    return matches[-1][0].strip()
+
+def read_bvids_from_txt(txt_path: str) -> List[str]:
+    try:
+        with open(txt_path, "r") as f:
+            bvids = [line.strip() for line in f if line.strip().startswith("BV")]
+        return bvids
+    except FileNotFoundError:
+        raise HTTPException(status_code=500, detail="Txt file not found")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error reading txt file: {str(e)}")
 
 async def scheduler():
     while True: