Update main.py

main.py

@@ -2,7 +2,7 @@ import os
 import asyncio
 from telethon.sync import TelegramClient
 from telethon.sessions import StringSession
-from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
 from fastapi.responses import FileResponse
 import uvicorn
 import tempfile
@@ -20,15 +20,12 @@ UPLOAD_DIR = tempfile.gettempdir()
 
 app = FastAPI()
 
-# NOTE: The in-memory file_map has been REMOVED for a stateless design.
-
 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:
-            # force_document=True ensures it's treated as a file with a consistent filename
             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.")
@@ -51,9 +48,9 @@ async def upload_file(file: UploadFile = File(...)):
         
         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.",
-            "channel": TARGET_CHANNEL,
             "telegram_details": {
                 "message_id": message.id,
                 "file_id": doc.id,
@@ -69,7 +66,7 @@ async def upload_file(file: UploadFile = File(...)):
             os.remove(file_location)
 
 @app.get("/v1/download/{message_id}")
-async def download_file(message_id: int):
+async def download_file(message_id: int, background_tasks: BackgroundTasks):
     """
     Downloads a file from Telegram using its unique message_id from the channel.
     This method is stateless and resilient to server restarts.
@@ -77,33 +74,24 @@ async def download_file(message_id: int):
     if not all([SESSION_STRING, API_ID, API_HASH]):
         raise HTTPException(status_code=500, detail="Server is not configured for downloads.")
 
-    download_path = None
     try:
         async with TelegramClient(StringSession(SESSION_STRING), API_ID, API_HASH) as client:
-            # Fetch the message directly using its ID
             message = await client.get_messages(TARGET_CHANNEL, ids=message_id)
             
             if not message or not message.media:
-                raise HTTPException(status_code=404, detail=f"A message with ID '{message_id}' was not found or contains no media.")
+                raise HTTPException(status_code=440, detail=f"A message with ID '{message_id}' was not found or contains no media.")
 
-            # Determine the original filename for the response
             original_filename = message.file.name if message.file and message.file.name else f"download_{message_id}"
-
-            # Download the file to a temporary location on the server
             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)
             
-            # Stream the file back to the client
             return FileResponse(download_path, media_type='application/octet-stream', filename=original_filename)
 
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to download from Telegram: {e}")
-    finally:
-        # Crucially, clean up the temporary file from the server after sending it
-        if download_path and os.path.exists(download_path):
-            os.remove(download_path)
-
-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)
+        # Catch potential errors like file not found in Telegram
+        raise HTTPException(status_code=500, detail=f"Failed to download from Telegram: {e}")