Update main.py

main.py

@@ -8,7 +8,6 @@ import os
 import sys
 import subprocess
 import json
-import tempfile
 import random
 import time
 import asyncio
@@ -18,10 +17,9 @@ 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, StreamingResponse
+from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File
+from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, HttpUrl
 import uvicorn
@@ -50,7 +48,7 @@ class BatchDownloadRequest(BaseModel):
     quality: str = "best"
     audio_only: bool = False
     use_cookies: Optional[bool] = None
-    max_concurrent: int = 2  # Limit concurrent downloads
+    max_concurrent: int = 2
 
 class VideoInfo(BaseModel):
     title: str
@@ -105,11 +103,16 @@ class HealthResponse(BaseModel):
     strategies_enabled: List[str]
     batch_support: bool
 
+class BatchFileListResponse(BaseModel):
+    batch_id: str
+    total_files: int
+    files: List[Dict[str, str]]
+
 # Initialize FastAPI app
 app = FastAPI(
     title="Batch YouTube Video Downloader",
     description="Download YouTube videos individually or in batches with cookie support",
-    version="4.0.0",
+    version="4.0.2",
     docs_url="/docs",
     redoc_url="/redoc"
 )
@@ -550,39 +553,6 @@ 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"""
@@ -723,7 +693,7 @@ async def read_root():
                 </div>
                 
                 <div class="feature new-feature">
-                    <strong>🆕 Bulk File Download:</strong> Download all files from a batch in one request
+                    <strong>🆕 Individual File Downloads:</strong> Download each file separately
                 </div>
                 
                 <div class="feature">
@@ -752,6 +722,7 @@ async def read_root():
                 <ul>
                     <li><code>POST /video/info</code> - Get single video information</li>
                     <li><code>POST /video/download</code> - Download single video</li>
+                    <li><code>GET /video/file/{filename}</code> - Download a specific file</li>
                 </ul>
                 
                 <h4>Batch Operations:</h4>
@@ -759,8 +730,7 @@ 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>
+                    <li><code>GET /batch/files/{batch_id}</code> - Get list of downloadable files</li>
                 </ul>
             </div>
             
@@ -873,7 +843,7 @@ async def health_check():
         "Enhanced Headers",
         "Concurrent Downloads",
         "Progress Tracking",
-        "Bulk File Download"
+        "Individual File Downloads"
     ]
     
     return HealthResponse(
@@ -1060,85 +1030,35 @@ 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}"
+@app.get("/batch/files/{batch_id}", response_model=BatchFileListResponse)
+async def get_batch_files(batch_id: str):
+    """Get list of downloadable files for a batch"""
+    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']}"
         )
-        
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.error(f"Error downloading all files: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+    
+    files = []
+    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):
+                files.append({
+                    "filename": os.path.basename(file_path),
+                    "url": f"/video/file/{os.path.basename(file_path)}"
+                })
+    
+    return BatchFileListResponse(
+        batch_id=batch_id,
+        total_files=len(files),
+        files=files
+    )
 
 @app.get("/video/file/{filename}")
 async def download_file(filename: str):
@@ -1180,5 +1100,4 @@ if __name__ == "__main__":
         port=port,
         reload=False,
         log_level="info"
-    )
-
+    )