Update main.py

main.py

@@ -29,12 +29,16 @@ async def log(message: str):
     if len(log_buffer) > 100:
         log_buffer.pop(0)
 
-def get_direct_link(file_id: str):
+def extract_file_id(url: str):
+    """Cleanly extracts ID from various URL formats."""
+    url = url.strip()
+    if "/v/" in url: return url.split("/v/")[1].split("/")[0]
+    if "/e/" in url: return url.split("/e/")[1].split("/")[0]
+    return url
+
+def get_download_info(file_id: str):
     """
-    Robust link extractor:
-    1. Handles 'Wait' errors (403) from Step 1.
-    2. Handles Direct URL returned immediately in Step 1.
-    3. Handles Ticket -> Wait -> Step 2 (JSON) -> Final URL.
+    Returns a tuple: (direct_url, filename)
     """
     base_url = "https://api.streamtape.com/file/dlticket"
     params = {'file': file_id, 'login': API_LOGIN, 'key': API_KEY}
@@ -42,7 +46,7 @@ def get_direct_link(file_id: str):
     max_retries = 5
     for attempt in range(max_retries):
         try:
-            print(f"Requesting ticket (Attempt {attempt+1})...")
+            print(f"Requesting ticket for {file_id} (Attempt {attempt+1})...")
             response = requests.get(base_url, params=params).json()
             
             status = response.get('status')
@@ -51,135 +55,120 @@ def get_direct_link(file_id: str):
 
             # --- CASE 1: SUCCESS (200 OK) ---
             if status == 200:
-                # Sub-case A: Direct URL provided immediately in Step 1
-                if result and 'url' in result and result['url']:
-                    print("✅ Direct URL found immediately in Step 1.")
-                    return result['url']
+                # Sub-case A: Direct URL provided immediately
+                if result.get('url'):
+                    # Name usually provided in result as well
+                    name = result.get('name', f"{file_id}.mp4")
+                    return result['url'], name
                 
-                # Sub-case B: Ticket provided (Standard Flow)
-                elif result and 'ticket' in result:
+                # Sub-case B: Ticket provided
+                elif result.get('ticket'):
                     ticket = result['ticket']
                     wait_time = result.get('wait_time', 5)
-                    print(f"🎟️ Ticket received. Waiting {wait_time}s...")
-                    
+                    print(f"🎟️ Ticket found. Waiting {wait_time}s...")
                     time.sleep(wait_time)
                     
-                    # Step 2: Request the download link
-                    # We treat this as a JSON request now, based on your findings
+                    # Step 2: Get Final Link & Name
                     dl_url = "https://api.streamtape.com/file/dl"
                     dl_params = {'file': file_id, 'ticket': ticket}
                     
-                    print("🔄 Requesting Step 2 (Final Link)...")
-                    dl_response = requests.get(dl_url, params=dl_params).json()
+                    dl_resp = requests.get(dl_url, params=dl_params).json()
                     
-                    if dl_response.get('status') == 200:
-                         final_url = dl_response['result']['url']
-                         print(f"✅ Final Link Found: {final_url}")
-                         return final_url
+                    if dl_resp.get('status') == 200:
+                         final_url = dl_resp['result']['url']
+                         final_name = dl_resp['result']['name']
+                         return final_url, final_name
                     else:
-                        raise Exception(f"Step 2 Error: {dl_response.get('msg')}")
-                
+                        raise Exception(f"Step 2 Error: {dl_resp.get('msg')}")
                 else:
-                    raise Exception("API returned 200 but neither URL nor Ticket was found.")
+                    raise Exception("API returned 200 but missing URL/Ticket.")
 
             # --- CASE 2: WAIT (403) ---
             elif status == 403 and "wait" in msg.lower():
                 wait_match = re.search(r'(\d+)', msg)
                 wait_seconds = int(wait_match.group(1)) if wait_match else 5
                 print(f"Rate Limit: Sleeping {wait_seconds}s...")
-                time.sleep(wait_seconds + 2) # +2s buffer
+                time.sleep(wait_seconds + 2)
                 continue
-                
             else:
                 raise Exception(f"API Error: {msg}")
 
         except Exception as e:
-            print(f"Error in attempt {attempt}: {e}")
-            if attempt == max_retries - 1:
-                raise e
+            print(f"Error: {e}")
+            if attempt == max_retries - 1: raise e
             time.sleep(2)
 
     raise Exception("Max retries exceeded")
 
-async def process_video_task(file_id: str, filename: str, hf_token: str = None, hf_repo: str = None):
-    local_path = os.path.join(DOWNLOAD_DIR, filename)
+async def batch_processor(urls: list, hf_token: str, hf_repo: str, repo_type: str):
+    """
+    Processes a list of URLs sequentially:
+    Download -> Upload -> Delete -> Next
+    """
+    await log(f"📦 Batch started with {len(urls)} links.")
     
-    # --- BROWSER HEADERS (Crucial for bypassing some blocks) ---
     headers = {
-        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36",
+        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) Chrome/115.0.0.0 Safari/537.36",
         "Referer": "https://streamtape.com/",
-        "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8",
-        "Accept-Language": "en-US,en;q=0.5",
         "Connection": "keep-alive"
     }
 
-    try:
-        await log(f"🚀 Starting task for: {filename}")
+    for index, url in enumerate(urls):
+        if not url.strip(): continue
         
-        # --- Step 1: Get Link ---
-        await log(f"Requesting link for ID {file_id}...")
-        try:
-            # Pass headers to the link extractor too just in case
-            direct_url = await asyncio.to_thread(get_direct_link, file_id)
-            await log(f"🔗 Link acquired successfully.")
-        except Exception as e:
-            await log(f"❌ Error getting link: {str(e)}")
-            return
-
-        # --- Step 2: Download with Session & Headers ---
-        await log(f"⬇️  Downloading from Streamtape...")
+        file_id = extract_file_id(url)
+        await log(f"▶️ Processing ({index+1}/{len(urls)}): ID {file_id}")
         
-        def download_chunked():
-            # Use a Session to manage connection pooling better
-            with requests.Session() as s:
-                s.headers.update(headers)
-                
-                # Timeout set to 15s to prevent hanging
-                with s.get(direct_url, stream=True, timeout=15) as r:
-                    r.raise_for_status()
-                    
-                    total_size = int(r.headers.get('content-length', 0))
-                    downloaded = 0
-                    
-                    with open(local_path, 'wb') as f:
-                        for chunk in r.iter_content(chunk_size=8192):
-                            if chunk:
-                                f.write(chunk)
-                                downloaded += len(chunk)
-            return downloaded
-
+        local_path = None
         try:
-            await asyncio.to_thread(download_chunked)
-            await log(f"✅ Download complete: {filename}")
-        except requests.exceptions.ConnectionError as e:
-            if "101" in str(e) or "unreachable" in str(e):
-                await log(f"⛔ BLOCKED: Hugging Face Free Tier firewall blocked the storage server.")
-                await log(f"💡 Solution: This code works perfectly on 'Local PC' or 'Paid Spaces'.")
-                return
-            else:
-                raise e
+            # 1. Get Link & Name
+            direct_url, filename = await asyncio.to_thread(get_download_info, file_id)
+            await log(f"   🔹 Name Detected: {filename}")
+            
+            local_path = os.path.join(DOWNLOAD_DIR, filename)
 
-        # --- Step 3: Upload ---
-        if hf_token and hf_repo:
-            await log(f"⬆️  Uploading to HF Repo: {hf_repo}...")
+            # 2. Download
+            await log(f"   ⬇️ Downloading...")
             
-            def upload_to_hf():
-                api = HfApi(token=hf_token)
-                api.upload_file(
-                    path_or_fileobj=local_path,
-                    path_in_repo=filename,
-                    repo_id=hf_repo,
-                    repo_type="dataset"
-                )
+            def download_file():
+                with requests.Session() as s:
+                    s.headers.update(headers)
+                    with s.get(direct_url, stream=True, timeout=45) as r:
+                        r.raise_for_status()
+                        with open(local_path, 'wb') as f:
+                            shutil.copyfileobj(r.raw, f)
             
-            try:
-                await asyncio.to_thread(upload_to_hf)
-                await log(f"🎉 Upload Success!")
-            except Exception as e:
-                await log(f"❌ Upload Failed: {str(e)}")
+            await asyncio.to_thread(download_file)
+            await log(f"   ✅ Downloaded to storage.")
+
+            # 3. Upload to HF
+            if hf_token and hf_repo:
+                await log(f"   ⬆️ Uploading to HF ({repo_type})...")
+                
+                def upload_file():
+                    api = HfApi(token=hf_token)
+                    api.upload_file(
+                        path_or_fileobj=local_path,
+                        path_in_repo=filename,
+                        repo_id=hf_repo,
+                        repo_type=repo_type
+                    )
+                
+                await asyncio.to_thread(upload_file)
+                await log(f"   🎉 Upload Complete!")
+                
+                # 4. DELETE LOCAL FILE (Cleanup)
+                if os.path.exists(local_path):
+                    os.remove(local_path)
+                    await log(f"   🗑️ Local file deleted to free space.")
+
+        except Exception as e:
+            await log(f"   ❌ Failed: {str(e)}")
+            # Even if failed, try to cleanup partial files
+            if local_path and os.path.exists(local_path):
+                os.remove(local_path)
 
-    except Exception as e:
-        await log(f"❌ Critical Error: {str(e)}")
+    await log("🏁 Batch Processing Completed.")
 
 # --- API ENDPOINTS ---
 
@@ -201,25 +190,30 @@ async def read_root(request: Request):
     files = os.listdir(DOWNLOAD_DIR) if os.path.exists(DOWNLOAD_DIR) else []
     return templates.TemplateResponse("index.html", {"request": request, "files": files})
 
-@app.post("/download")
-async def start_download(
+@app.post("/start_batch")
+async def start_batch(
     request: Request, 
     background_tasks: BackgroundTasks,
-    url: str = Form(...), 
-    filename: str = Form(...),
+    urls: str = Form(...), 
     hf_token: str = Form(None),
-    hf_repo: str = Form(None)
+    hf_repo: str = Form(None),
+    repo_type: str = Form("dataset")
 ):
-    try:
-        # Robust ID Extraction
-        if "/v/" in url: file_id = url.split("/v/")[1].split("/")[0]
-        elif "/e/" in url: file_id = url.split("/e/")[1].split("/")[0]
-        else: file_id = url.strip()
-        
-        background_tasks.add_task(process_video_task, file_id, filename, hf_token, hf_repo)
-        msg = f"Task started for {filename}"
-    except Exception as e:
-        msg = f"Error: {str(e)}"
-
-    files = os.listdir(DOWNLOAD_DIR) if os.path.exists(DOWNLOAD_DIR) else []
-    return templates.TemplateResponse("index.html", {"request": request, "files": files, "message": msg})
+    # Split text area by newlines
+    url_list = [line.strip() for line in urls.splitlines() if line.strip()]
+    
+    if not url_list:
+        return templates.TemplateResponse("index.html", {
+            "request": request, 
+            "files": os.listdir(DOWNLOAD_DIR), 
+            "message": "❌ No URLs provided."
+        })
+
+    # Start background process
+    background_tasks.add_task(batch_processor, url_list, hf_token, hf_repo, repo_type)
+    
+    return templates.TemplateResponse("index.html", {
+        "request": request, 
+        "files": os.listdir(DOWNLOAD_DIR), 
+        "message": f"🚀 Batch started for {len(url_list)} files."
+    })