From ab15aff021d85b25fbd558d134fbf6a9dd6cd6ea Mon Sep 17 00:00:00 2001
From: Abi Raja <abimanyuraja@gmail.com>
Date: Tue, 30 Jan 2024 14:32:57 -0500
Subject: [PATCH] Support screenshot by URL for all paying customers

---
 backend/config.py                           |  3 ++
 backend/routes/screenshot.py                | 31 +++++++++++++++++++--
 frontend/src/components/UrlInputSection.tsx | 23 +++++++++------
 3 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/backend/config.py b/backend/config.py
index aa6ca22..f944539 100644
--- a/backend/config.py
+++ b/backend/config.py
@@ -11,4 +11,7 @@ SHOULD_MOCK_AI_RESPONSE = bool(os.environ.get("MOCK", False))
 IS_PROD = os.environ.get("IS_PROD", False)
 
 # Hosted version only
+
+PLATFORM_SCREENSHOTONE_API_KEY = os.environ.get("PLATFORM_SCREENSHOTONE_API_KEY", "")
+
 BACKEND_SAAS_URL = os.environ.get("BACKEND_SAAS_URL", "")
diff --git a/backend/routes/screenshot.py b/backend/routes/screenshot.py
index 258cd7e..0ce497e 100644
--- a/backend/routes/screenshot.py
+++ b/backend/routes/screenshot.py
@@ -2,6 +2,9 @@ import base64
 from fastapi import APIRouter
 from pydantic import BaseModel
 import httpx
+from config import PLATFORM_SCREENSHOTONE_API_KEY
+
+from routes.saas_utils import does_user_have_subscription_credits
 
 router = APIRouter()
 
@@ -12,10 +15,30 @@ def bytes_to_data_url(image_bytes: bytes, mime_type: str) -> str:
 
 
 async def capture_screenshot(
-    target_url: str, api_key: str, device: str = "desktop"
+    target_url: str, api_key: str | None, auth_token: str, device: str = "desktop"
 ) -> bytes:
     api_base_url = "https://api.screenshotone.com/take"
 
+    # Get auth token
+    if not auth_token:
+        raise Exception("No auth token with capture_screenshot")
+
+    # If API key is not passed in, only use the platform ScreenshotOne API key if the user is a subscriber
+    if not api_key:
+        res = await does_user_have_subscription_credits(auth_token)
+        if res.status == "not_subscriber":
+            raise Exception(
+                "capture_screenshot - User is not subscriber and has no API key"
+            )
+        elif res.status == "subscriber_has_credits":
+            api_key = PLATFORM_SCREENSHOTONE_API_KEY
+        elif res.status == "subscriber_has_no_credits":
+            raise Exception("capture_screenshot - User has no credits")
+        else:
+            raise Exception(
+                "capture_screenshot - Unknown error occurred when checking subscription credits"
+            )
+
     params = {
         "access_key": api_key,
         "url": target_url,
@@ -44,7 +67,8 @@ async def capture_screenshot(
 
 class ScreenshotRequest(BaseModel):
     url: str
-    apiKey: str
+    apiKey: str | None
+    authToken: str
 
 
 class ScreenshotResponse(BaseModel):
@@ -56,9 +80,10 @@ async def app_screenshot(request: ScreenshotRequest):
     # Extract the URL from the request body
     url = request.url
     api_key = request.apiKey
+    auth_token = request.authToken
 
     # TODO: Add error handling
-    image_bytes = await capture_screenshot(url, api_key=api_key)
+    image_bytes = await capture_screenshot(url, api_key=api_key, auth_token=auth_token)
 
     # Convert the image bytes to a data url
     data_url = bytes_to_data_url(image_bytes, "image/png")
diff --git a/frontend/src/components/UrlInputSection.tsx b/frontend/src/components/UrlInputSection.tsx
index e2137e1..2a5e07e 100644
--- a/frontend/src/components/UrlInputSection.tsx
+++ b/frontend/src/components/UrlInputSection.tsx
@@ -3,6 +3,8 @@ import { HTTP_BACKEND_URL } from "../config";
 import { Button } from "./ui/button";
 import { Input } from "./ui/input";
 import { toast } from "react-hot-toast";
+import { useStore } from "../store/store";
+import { useAuth } from "@clerk/clerk-react";
 
 interface Props {
   screenshotOneApiKey: string | null;
@@ -13,28 +15,31 @@ export function UrlInputSection({ doCreate, screenshotOneApiKey }: Props) {
   const [isLoading, setIsLoading] = useState(false);
   const [referenceUrl, setReferenceUrl] = useState("");
 
+  // Hosted version only
+  const subscriberTier = useStore((state) => state.subscriberTier);
+  const { getToken } = useAuth();
+
   async function takeScreenshot() {
-    if (!screenshotOneApiKey) {
-      toast.error(
-        "Please add a ScreenshotOne API key in the Settings dialog. This is optional - you can also drag/drop and upload images directly.",
-        { duration: 8000 }
-      );
-      return;
+    if (!referenceUrl) {
+      return toast.error("Please enter a URL");
     }
 
-    if (!referenceUrl) {
-      toast.error("Please enter a URL");
-      return;
+    if (!screenshotOneApiKey && subscriberTier === "free") {
+      return toast.error(
+        "Please upgrade to a paid plan to use the screenshot feature."
+      );
     }
 
     if (referenceUrl) {
       try {
         setIsLoading(true);
+        const authToken = await getToken();
         const response = await fetch(`${HTTP_BACKEND_URL}/api/screenshot`, {
           method: "POST",
           body: JSON.stringify({
             url: referenceUrl,
             apiKey: screenshotOneApiKey,
+            authToken,
           }),
           headers: {
             "Content-Type": "application/json",