updated

App/Youtube/youtubeRouter.py

@@ -1,10 +1,11 @@
 from fastapi import APIRouter, HTTPException, BackgroundTasks
-from subprocess import run
 import tempfile
 import os
 import logging
 import aiohttp
 import asyncio
+import aiofiles
+import subprocess
 from .Schema import YouTubeUploadTask
 
 # Configure logging
@@ -15,28 +16,81 @@ router = APIRouter(prefix="/youtube", tags=["youtube"])
 
 
 async def download_file(url: str) -> str:
-    """Download file from URL and save to temp directory."""
+    """Download file from URL and save to temp directory with correct extension based on metadata."""
     try:
         async with aiohttp.ClientSession() as session:
             async with session.get(url) as response:
-                if response.status != 200:
+                if response.status not in (200, 206):
                     raise HTTPException(
                         status_code=400,
                         detail=f"Failed to download file: HTTP {response.status}",
                     )
 
-                # Create temp file with appropriate extension
-                content_type = response.headers.get("content-type", "")
-                ext = ".mp4" if "video" in content_type else ".tmp"
-
-                fd, temp_path = tempfile.mkstemp(suffix=ext)
+                # Create temp file with .tmp extension
+                fd, temp_path = tempfile.mkstemp(suffix=".tmp")
 
                 # Write content to temp file
-                with os.fdopen(fd, "wb") as tmp:
+                async with aiofiles.open(fd, "wb") as tmp:
                     while chunk := await response.content.read(8192):
-                        tmp.write(chunk)
-
-                return temp_path
+                        await tmp.write(chunk)
+
+                # Determine file format using ffprobe
+                try:
+                    # Get codec information
+                    codec_cmd = [
+                        "ffprobe",
+                        "-v",
+                        "error",
+                        "-select_streams",
+                        "v:0",
+                        "-show_entries",
+                        "stream=codec_name",
+                        "-of",
+                        "default=noprint_wrappers=1:nokey=1",
+                        temp_path,
+                    ]
+                    codec_result = subprocess.run(
+                        codec_cmd, capture_output=True, text=True
+                    )
+                    codec = codec_result.stdout.strip().lower()
+
+                    # Get container format
+                    format_cmd = [
+                        "ffprobe",
+                        "-v",
+                        "error",
+                        "-show_entries",
+                        "format=format_name",
+                        "-of",
+                        "default=noprint_wrappers=1:nokey=1",
+                        temp_path,
+                    ]
+                    format_result = subprocess.run(
+                        format_cmd, capture_output=True, text=True
+                    )
+                    format_name = format_result.stdout.strip().lower()
+
+                    # Determine extension based on codec and container
+                    if "mp4" in format_name or codec in ("h264", "aac"):
+                        ext = ".mp4"
+                    elif "webm" in format_name or codec == "vp9":
+                        ext = ".webm"
+                    elif "avi" in format_name:
+                        ext = ".avi"
+                    else:
+                        # Fallback to content-type if metadata unavailable
+                        content_type = response.headers.get("content-type", "")
+                        ext = ".mp4" if "video" in content_type else ".tmp"
+                except Exception:
+                    # Fallback if ffprobe fails
+                    content_type = response.headers.get("content-type", "")
+                    ext = ".mp4" if "video" in content_type else ".tmp"
+
+                # Rename file with correct extension
+                new_temp_path = f"{temp_path[:-4]}{ext}"
+                os.rename(temp_path, new_temp_path)
+
+                return new_temp_path
     except Exception as e:
         logger.error(f"Error downloading file: {str(e)}")
         raise HTTPException(status_code=500, detail=f"File download failed: {str(e)}")
@@ -71,20 +125,24 @@ async def upload_video_background(task: YouTubeUploadTask):
             command.extend(["-thumbnail", task.thumbnail])
 
         logger.info("Executing upload command")
-        result = run(command, capture_output=True, text=True)
+        process = await asyncio.create_subprocess_exec(
+            *command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        stdout, stderr = await process.communicate()
 
-        if result.returncode != 0:
-            logger.error(f"Upload failed: {result.stderr}")
-            raise Exception(f"Upload failed: {result.stderr}")
+        if process.returncode != 0:
+            logger.error(f"Upload failed: {stderr.decode()}")
+            raise Exception(f"Upload failed: {stderr.decode()}")
 
         logger.info("Upload completed successfully")
-        logger.debug(f"Upload output: {result.stdout}")
+        logger.debug(f"Upload output: {stdout.decode()}")
 
     except Exception as e:
         logger.error(f"Error in upload process: {str(e)}")
         raise
     finally:
-        # Clean up temp file
         if temp_file and os.path.exists(temp_file):
             try:
                 os.remove(temp_file)
@@ -104,7 +162,6 @@ async def upload_video_to_youtube(
     try:
         background_tasks.add_task(upload_video_background, task)
         return {"message": "Upload task started", "status": "processing"}
-
     except Exception as e:
         logger.error(f"Error initiating upload task: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))