Update main.py

main.py

@@ -2,8 +2,9 @@ import os
 import asyncio
 from telethon.sync import TelegramClient
 from telethon.sessions import StringSession
-from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
-from fastapi.responses import FileResponse
+from telethon.tl.functions.upload import GetFileRequest
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.responses import RedirectResponse
 import uvicorn
 import tempfile
 
@@ -15,24 +16,11 @@ API_HASH = os.environ.get('API_HASH')
 # The Telegram channel to use as file storage
 TARGET_CHANNEL = 'https://t.me/TestNAI01'
 
-# Use the system's temporary directory for transient files
+# Use the system's temporary directory for transient uploads
 UPLOAD_DIR = tempfile.gettempdir()
 
 app = FastAPI()
 
-async def upload_to_telegram(file_path: str):
-    """Uploads a file and returns the full Telegram message object."""
-    if not all([SESSION_STRING, API_ID, API_HASH]):
-        raise ValueError("Server is not configured with Telegram credentials.")
-    async with TelegramClient(StringSession(SESSION_STRING), API_ID, API_HASH) as client:
-        try:
-            message = await client.send_file(TARGET_CHANNEL, file_path, force_document=True)
-            if not message or not message.media:
-                raise HTTPException(status_code=500, detail="Upload failed to return a valid message.")
-            return message
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Failed to upload to Telegram: {e}")
-
 @app.post("/v1/upload")
 async def upload_file(file: UploadFile = File(...)):
     """
@@ -43,55 +31,72 @@ async def upload_file(file: UploadFile = File(...)):
         with open(file_location, "wb+") as file_object:
             file_object.write(file.file.read())
 
-        message = await upload_to_telegram(file_location)
-        doc = message.document
-        
-        original_filename = next((attr.file_name for attr in doc.attributes if hasattr(attr, 'file_name')), file.filename)
+        async with TelegramClient(StringSession(SESSION_STRING), API_ID, API_HASH) as client:
+            message = await client.send_file(TARGET_CHANNEL, file_location, force_document=True)
+            if not message or not message.media:
+                raise HTTPException(status_code=500, detail="Upload failed to return a valid message.")
+            
+            doc = message.document
+            original_filename = next((attr.file_name for attr in doc.attributes if hasattr(attr, 'file_name')), file.filename)
 
-        # Return details without the channel link, as requested.
-        return {
-            "message": "File uploaded successfully. Use the 'message_id' to download.",
-            "telegram_details": {
-                "message_id": message.id,
-                "file_id": doc.id,
-                "file_reference": doc.file_reference.hex() if doc.file_reference else None,
-                "filename": original_filename,
-                "size_bytes": doc.size,
-                "mime_type": doc.mime_type,
-                "date": message.date.isoformat()
+            return {
+                "message": "File uploaded successfully. Use the 'message_id' to download.",
+                "telegram_details": {
+                    "message_id": message.id,
+                    "filename": original_filename,
+                    "size_bytes": doc.size,
+                    "mime_type": doc.mime_type,
+                    "date": message.date.isoformat()
+                }
             }
-        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to upload to Telegram: {e}")
     finally:
         if os.path.exists(file_location):
             os.remove(file_location)
 
 @app.get("/v1/download/{message_id}")
-async def download_file(message_id: int, background_tasks: BackgroundTasks):
+async def download_file(message_id: int):
     """
-    Downloads a file from Telegram using its unique message_id from the channel.
-    This method is stateless and resilient to server restarts.
+    Generates a direct, temporary download link from Telegram's CDN and redirects the client to it.
+    This is extremely fast and efficient.
     """
     if not all([SESSION_STRING, API_ID, API_HASH]):
         raise HTTPException(status_code=500, detail="Server is not configured for downloads.")
 
     try:
         async with TelegramClient(StringSession(SESSION_STRING), API_ID, API_HASH) as client:
+            # Get the message object
             message = await client.get_messages(TARGET_CHANNEL, ids=message_id)
             
-            if not message or not message.media:
-                raise HTTPException(status_code=440, detail=f"A message with ID '{message_id}' was not found or contains no media.")
+            if not message or not message.document:
+                raise HTTPException(status_code=404, detail=f"A message with ID '{message_id}' was not found or is not a file.")
 
-            original_filename = message.file.name if message.file and message.file.name else f"download_{message_id}"
-            download_path = await client.download_media(message, file=UPLOAD_DIR)
-
-            # **THE FIX IS HERE:**
-            # The os.remove is now a background task. It will only run AFTER the
-            # FileResponse has been fully sent to the client, preventing the error.
-            if download_path:
-                background_tasks.add_task(os.remove, download_path)
+            # Get the data center information where the file is stored
+            dc_id, location = message.document.dc_id, message.document
+            
+            # Use a low-level request to generate a temporary download URL
+            # This is the key to creating a CDN link
+            req = GetFileRequest(
+                location=location,
+                cdn_supported=True,
+                offset=0,
+                limit=0
+            )
+            result = await client(req)
             
-            return FileResponse(download_path, media_type='application/octet-stream', filename=original_filename)
+            # The result contains the temporary URL.
+            # Format: https://{cdn_redirect_ip}/.../{file_token}
+            cdn_url = f"https://{result.cdn_redirect}/file/{result.bytes.hex()}"
+            
+            # Immediately redirect the user's client to this URL
+            return RedirectResponse(url=cdn_url, status_code=307)
 
     except Exception as e:
-        # Catch potential errors like file not found in Telegram
-        raise HTTPException(status_code=500, detail=f"Failed to download from Telegram: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to generate download link: {e}")
+
+if __name__ == "__main__":
+    if not all([SESSION_STRING, API_ID, API_HASH]):
+        print("FATAL ERROR: The environment variables SESSION_STRING, API_ID, and API_HASH must be set.")
+    else:
+        uvicorn.run(app, host="0.0.0.0", port=8000)