Update main.py

main.py

@@ -18,9 +18,10 @@ from pathlib import Path
 from typing import Optional, Dict, Any, List
 from datetime import datetime
 from concurrent.futures import ThreadPoolExecutor
+import io
 
 from fastapi import FastAPI, HTTPException, BackgroundTasks, Request, UploadFile, File, Form
-from fastapi.responses import FileResponse, HTMLResponse
+from fastapi.responses import FileResponse, HTMLResponse, StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, HttpUrl
 import uvicorn
@@ -549,6 +550,39 @@ class BatchYouTubeDownloader:
 # Global downloader instance
 downloader = BatchYouTubeDownloader()
 
+def generate_multipart_response(file_paths: List[str], boundary: str = None):
+    """Generate a multipart response with multiple files"""
+    if boundary is None:
+        boundary = f"----formdata-{uuid.uuid4().hex}"
+    
+    def file_generator():
+        for file_path in file_paths:
+            if os.path.exists(file_path):
+                filename = os.path.basename(file_path)
+                file_size = os.path.getsize(file_path)
+                
+                # Write multipart boundary and headers
+                yield f"--{boundary}\r\n".encode()
+                yield f"Content-Disposition: form-data; name=\"file\"; filename=\"{filename}\"\r\n".encode()
+                yield f"Content-Type: application/octet-stream\r\n".encode()
+                yield f"Content-Length: {file_size}\r\n".encode()
+                yield "\r\n".encode()
+                
+                # Stream file content
+                with open(file_path, 'rb') as f:
+                    while True:
+                        chunk = f.read(8192)
+                        if not chunk:
+                            break
+                        yield chunk
+                
+                yield "\r\n".encode()
+        
+        # Final boundary
+        yield f"--{boundary}--\r\n".encode()
+    
+    return file_generator(), boundary
+
 @app.get("/", response_class=HTMLResponse)
 async def read_root():
     """Serve the main HTML interface with batch support and cookie upload"""
@@ -688,6 +722,10 @@ async def read_root():
                     <strong>🆕 Progress Tracking:</strong> Monitor batch download progress in real-time
                 </div>
                 
+                <div class="feature new-feature">
+                    <strong>🆕 Bulk File Download:</strong> Download all files from a batch in one request
+                </div>
+                
                 <div class="feature">
                     <strong>🍪 Cookie Support:</strong> Upload cookies for better success rates
                 </div>
@@ -721,6 +759,8 @@ async def read_root():
                     <li><code>POST /batch/info</code> - Get info for multiple videos</li>
                     <li><code>POST /batch/download</code> - Start batch download</li>
                     <li><code>GET /batch/status/{batch_id}</code> - Check batch progress</li>
+                    <li><code>GET /batch/download-all/{batch_id}</code> - Download all files from batch</li>
+                    <li><code>GET /download-all</code> - Download ALL files from server</li>
                 </ul>
             </div>
             
@@ -832,7 +872,8 @@ async def health_check():
         "Smart Retry Logic",
         "Enhanced Headers",
         "Concurrent Downloads",
-        "Progress Tracking"
+        "Progress Tracking",
+        "Bulk File Download"
     ]
     
     return HealthResponse(
@@ -1019,6 +1060,86 @@ async def get_batch_status(batch_id: str):
     status = batch_status_store[batch_id]
     return BatchStatus(**status)
 
+@app.get("/batch/download-all/{batch_id}")
+async def download_all_batch_files(batch_id: str):
+    """Download all files from a completed batch as multipart response"""
+    try:
+        if batch_id not in batch_status_store:
+            raise HTTPException(status_code=404, detail="Batch not found")
+        
+        batch_status = batch_status_store[batch_id]
+        
+        if batch_status["status"] != "completed":
+            raise HTTPException(
+                status_code=400, 
+                detail=f"Batch is not completed yet. Current status: {batch_status['status']}"
+            )
+        
+        # Get all successfully downloaded files
+        file_paths = []
+        for result in batch_status["results"]:
+            if result.get("success") and result.get("download_path"):
+                file_path = result["download_path"]
+                if os.path.exists(file_path):
+                    file_paths.append(file_path)
+        
+        if not file_paths:
+            raise HTTPException(status_code=404, detail="No files found for this batch")
+        
+        # Generate multipart response
+        file_generator, boundary = generate_multipart_response(file_paths)
+        
+        headers = {
+            "Content-Type": f"multipart/form-data; boundary={boundary}",
+            "Content-Disposition": f"attachment; filename=\"batch_{batch_id}_files.multipart\""
+        }
+        
+        return StreamingResponse(
+            file_generator,
+            headers=headers,
+            media_type=f"multipart/form-data; boundary={boundary}"
+        )
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error downloading batch files: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/download-all")
+async def download_all_files():
+    """Download all files from the server as multipart response"""
+    try:
+        # Get all files from the download directory
+        download_dir = Path(downloader.download_dir)
+        all_files = [f for f in download_dir.glob("*") if f.is_file()]
+        
+        if not all_files:
+            raise HTTPException(status_code=404, detail="No files found on server")
+        
+        # Convert to string paths
+        file_paths = [str(f) for f in all_files]
+        
+        # Generate multipart response
+        file_generator, boundary = generate_multipart_response(file_paths)
+        
+        headers = {
+            "Content-Type": f"multipart/form-data; boundary={boundary}",
+            "Content-Disposition": f"attachment; filename=\"all_server_files.multipart\""
+        }
+        
+        return StreamingResponse(
+            file_generator,
+            headers=headers,
+            media_type=f"multipart/form-data; boundary={boundary}"
+        )
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error downloading all files: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
 @app.get("/video/file/{filename}")
 async def download_file(filename: str):
     """Serve downloaded files"""