Update main.py

main.py

@@ -6,7 +6,6 @@ from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.responses import FileResponse
 import uvicorn
 import tempfile
-from typing import Dict
 
 # --- Configuration from environment variables ---
 SESSION_STRING = os.environ.get('SESSION_STRING')
@@ -21,10 +20,7 @@ UPLOAD_DIR = tempfile.gettempdir()
 
 app = FastAPI()
 
-# --- In-Memory Database for Filename to Message ID mapping ---
-# NOTE: This dictionary is not persistent. It will be cleared if the server restarts.
-# For production, replace this with a real database (e.g., Redis, SQLite).
-file_map: Dict[str, int] = {}
+# 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."""
@@ -32,6 +28,7 @@ async def upload_to_telegram(file_path: str):
         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.")
@@ -42,7 +39,7 @@ async def upload_to_telegram(file_path: str):
 @app.post("/v1/upload")
 async def upload_file(file: UploadFile = File(...)):
     """
-    Uploads a file to Telegram, stores its mapping, and returns full API details.
+    Uploads a file to Telegram and returns its persistent message_id and other details.
     """
     file_location = os.path.join(UPLOAD_DIR, file.filename)
     try:
@@ -52,15 +49,10 @@ async def upload_file(file: UploadFile = File(...)):
         message = await upload_to_telegram(file_location)
         doc = message.document
         
-        # Extract the original filename from the document attributes
         original_filename = next((attr.file_name for attr in doc.attributes if hasattr(attr, 'file_name')), file.filename)
 
-        # Store the mapping in our in-memory "database"
-        file_map[original_filename] = message.id
-
-        # Return a rich, detailed JSON response from the Telegram API
         return {
-            "message": "File uploaded successfully.",
+            "message": "File uploaded successfully. Use the 'message_id' to download.",
             "channel": TARGET_CHANNEL,
             "telegram_details": {
                 "message_id": message.id,
@@ -76,28 +68,37 @@ async def upload_file(file: UploadFile = File(...)):
         if os.path.exists(file_location):
             os.remove(file_location)
 
-@app.get("/v1/download/{filename}")
-async def download_file(filename: str):
+@app.get("/v1/download/{message_id}")
+async def download_file(message_id: int):
     """
-    Downloads a file from Telegram by its filename using the stored mapping.
+    Downloads a file from Telegram using its unique message_id from the channel.
+    This method is stateless and resilient to server restarts.
     """
-    # Look up the message_id from our in-memory map
-    message_id = file_map.get(filename)
-    if not message_id:
-        raise HTTPException(status_code=404, detail=f"File '{filename}' not found. The server may have restarted or the file was never uploaded.")
+    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="File found in map, but message is missing or has no media in Telegram.")
+                raise HTTPException(status_code=404, 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)
-            return FileResponse(download_path, media_type='application/octet-stream', filename=filename)
+            
+            # 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)