feat: implement 10s image caching by stream id in /app/{id}.jpg

main.py

@@ -2,7 +2,11 @@ import cv2
 import yt_dlp
 import numpy as np
 import json
+import time
+import base64
 
+_snapshot_cache = {}
+CACHE_TTL = 10
 
 def capture_stream_snapshot(youtube_url, output_file="snapshot.jpg", toBuffer=False):
     # 1. Configure yt-dlp to get the direct stream URL
@@ -29,7 +33,7 @@ def capture_stream_snapshot(youtube_url, output_file="snapshot.jpg", toBuffer=Fa
 
     if not cap.isOpened():
         print("Error: Could not open video stream.")
-        return None if toBuffer else None
+        return None
 
     # Read a single frame
     success, frame = cap.read()
@@ -39,138 +43,126 @@ def capture_stream_snapshot(youtube_url, output_file="snapshot.jpg", toBuffer=Fa
             # Encode frame as JPEG to memory buffer
             ret, buf = cv2.imencode(".jpg", frame)
             if ret:
-                print("Snapshot captured to buffer.")
                 cap.release()
                 return buf.tobytes()
             else:
-                print("Error: Could not encode frame to buffer.")
                 cap.release()
                 return None
         else:
             # 3. Save the frame as an image
             cv2.imwrite(output_file, frame)
-            print(f"Snapshot saved to {output_file}")
+            cap.release()
+            return True
     else:
         print("Error: Could not read frame from stream.")
-
-    # Cleanup
-    cap.release()
+        cap.release()
+        return None
 
 
 # Shsf Handler
 def main(args):
-    route = args.get("route")
+    route = args.get("route", "")
+    
+    # Handle /app/{id}.jpg
+    if route.startswith("app/") and route.endswith(".jpg"):
+        url_id = route.split("/")[-1].replace(".jpg", "")
+        now = time.time()
+        
+        if url_id in _snapshot_cache:
+            ts, encoded = _snapshot_cache[url_id]
+            if now - ts < CACHE_TTL:
+                return {
+                    "_shsf": "v1", # Changed to v1 for simpler binary response if v2 is picky
+                    "_code": 200,
+                    "_res": encoded,
+                    "_is_base64": True,
+                    "_headers": { "Content-Type": "image/jpeg", "X-Cache": "HIT" }
+                }
+        
+        # Cache miss
+        yt_url = f"https://www.youtube.com/watch?v={url_id}"
+        image_bytes = capture_stream_snapshot(yt_url, toBuffer=True)
+        if image_bytes:
+            encoded_content = base64.b64encode(image_bytes).decode('utf-8')
+            _snapshot_cache[url_id] = (now, encoded_content)
+            return {
+                "_shsf": "v1",
+                "_code": 200,
+                "_res": encoded_content,
+                "_is_base64": True,
+                "_headers": { "Content-Type": "image/jpeg", "X-Cache": "MISS" }
+            }
+
     if route == "snapshot":
         url = args.get("body", "")
         try:
             url = json.loads(url)
         except json.JSONDecodeError:
-            print("Error: Invalid JSON input.")
             return {
                 "_shsf": "v2",
                 "_code": 400,
                 "_res": {"message": "Invalid JSON input."},
-                "_headers": {
-                    "Content-Type": "application/json",
-                },
+                "_headers": { "Content-Type": "application/json" },
             }
         url = url.get("url", "")
         if not url:
-            print("Error: URL not provided.")
             return {
                 "_shsf": "v2",
                 "_code": 400,
                 "_res": {"message": "URL not provided."},
-                "_headers": {
-                    "Content-Type": "application/json",
-                },
+                "_headers": { "Content-Type": "application/json" },
             }
         
         try:
-            capture_stream_snapshot(url, output_file="/tmp/snapshot.jpg")
-            content = None
-            with open("/tmp/snapshot.jpg", "rb") as f:
-                content = f.read()
+            image_bytes = capture_stream_snapshot(url, toBuffer=True)
+            if image_bytes:
+                encoded_content = base64.b64encode(image_bytes).decode('utf-8')
+                return {
+                    "_shsf": "v2",
+                    "_code": 200,
+                    "_res": { "message": "done", "buffer": encoded_content },
+                    "_headers": { "Content-Type": "application/json" },
+                }
         except Exception as e:
-            print(f"Error: {e}")
             return {
                 "_shsf": "v2",
                 "_code": 500,
-                "_res": {"message": f"Could not capture snapshot: {str(e)}"},
-                "_headers": {
-                    "Content-Type": "application/json",
-                },
+                "_res": {"message": str(e)},
+                "_headers": { "Content-Type": "application/json" },
             }
         
-        if content is None:
-            print("Error: Could not capture snapshot.")
-            return {
-                "_shsf": "v2",
-                "_code": 500,
-                "_res": {"message": "Could not capture snapshot."},
-                "_headers": {
-                    "Content-Type": "application/json",
-                },
-            }
-        
-        # Convert content to base64 for JSON response
-        import base64
-        encoded_content = base64.b64encode(content).decode('utf-8')
-        
         return {
             "_shsf": "v2",
-            "_code": 200,
-            "_res": {
-                "message": "done",
-                "buffer": encoded_content
-            },
-            "_headers": {
-                "Content-Type": "application/json",
-            },
+            "_code": 500,
+            "_res": {"message": "Could not capture snapshot."},
+            "_headers": { "Content-Type": "application/json" },
         }
+
     elif route == "health":
         return {
             "_shsf": "v2",
             "_code": 200,
             "_res": {"message": "Service is healthy."},
-            "_headers": {
-                "Content-Type": "application/json",
-            },
+            "cache_size": len(_snapshot_cache),
+            "_headers": { "Content-Type": "application/json" },
         }
+    
     elif route == "default": # UI Route
         try:
-            content = ""
             with open("/app/ui.html", "r") as f:
                 content = f.read()
-        except FileNotFoundError:
-            print("Error: /app/ui.html not found.")
+            return {
+                "_shsf": "v2",
+                "_code": 200,
+                "_res": content,
+                "_headers": { "Content-Type": "text/html" },
+            }
+        except Exception as e:
             return {
                 "_shsf": "v2",
                 "_code": 404,
-                "_res": {"message": "UI file not found."},
-                "_headers": {
-                    "Content-Type": "application/json",
-                },
+                "_res": {"message": "UI file not found or error."},
             }
-        except Exception as e:
-            print(f"Error loading UI: {e}")
-            return {
-                "_shsf": "v2",
-                "_code": 500,
-                "_res": {"message": f"Server error: {str(e)}"},
-                "_headers": {
-                    "Content-Type": "application/json",
-                },
-            }
-        return {
-            "_shsf": "v2",
-            "_code": 200,
-            "_res": content,
-            "_headers": {
-                "Content-Type": "text/html",
-            },
-        }
         
 if __name__ == "__main__":
-    print("omg im on main")
-    capture_stream_snapshot("https://www.youtube.com/watch?v=n15V_fCsl_c", output_file="snapshot.jpg")
+    print("running stream-shot")