Update main.py

main.py

@@ -6,6 +6,7 @@ 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')
@@ -20,20 +21,20 @@ 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] = {}
+
 async def upload_to_telegram(file_path: str):
-    """
-    Uploads a file to the Telegram channel and returns the message object.
-    """
+    """Uploads a file and returns the full Telegram message object."""
     if not all([SESSION_STRING, API_ID, API_HASH]):
-        raise ValueError("SESSION_STRING, API_ID, and API_HASH must be set as environment variables.")
-
+        raise ValueError("Server is not configured with Telegram credentials.")
     async with TelegramClient(StringSession(SESSION_STRING), API_ID, API_HASH) as client:
         try:
-            # Send the file and wait for the message object to be returned
-            message = await client.send_file(TARGET_CHANNEL, file_path)
-            # Check if the upload was successful and returned a message with media
+            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="File uploaded but failed to receive a valid message object from Telegram.")
+                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}")
@@ -41,62 +42,62 @@ async def upload_to_telegram(file_path: str):
 @app.post("/v1/upload")
 async def upload_file(file: UploadFile = File(...)):
     """
-    Receives a file, uploads it to Telegram, and returns the Telegram message ID.
+    Uploads a file to Telegram, stores its mapping, and returns full API details.
     """
     file_location = os.path.join(UPLOAD_DIR, file.filename)
     try:
         with open(file_location, "wb+") as file_object:
             file_object.write(file.file.read())
 
-        # Upload the file and get the message object
         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)
 
-        # The message ID is the permanent identifier for the file
-        telegram_file_id = message.id
-        original_filename = message.file.name if message.file and message.file.name else 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 to Telegram.",
-            "telegram_file_id": telegram_file_id,
-            "filename": original_filename,
-            "channel": TARGET_CHANNEL
+            "message": "File uploaded successfully.",
+            "channel": TARGET_CHANNEL,
+            "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()
+            }
         }
     finally:
-        # Clean up the local temporary file
         if os.path.exists(file_location):
             os.remove(file_location)
 
-@app.get("/v1/download/{message_id}")
-async def download_file(message_id: int):
+@app.get("/v1/download/{filename}")
+async def download_file(filename: str):
     """
-    Downloads a file from Telegram using its message ID.
+    Downloads a file from Telegram by its filename using the stored mapping.
     """
-    if not all([SESSION_STRING, API_ID, API_HASH]):
-        raise HTTPException(status_code=500, detail="Server is not configured with Telegram credentials.")
+    # 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.")
 
     download_path = None
     try:
         async with TelegramClient(StringSession(SESSION_STRING), API_ID, API_HASH) as client:
-            # Get the specific message from the channel using the message_id
             message = await client.get_messages(TARGET_CHANNEL, ids=message_id)
-
             if not message or not message.media:
-                raise HTTPException(status_code=404, detail="Message not found or it contains no media.")
+                raise HTTPException(status_code=404, detail="File found in map, but message is missing or has no media in Telegram.")
 
-            # Download the media to the temporary directory
-            # Telethon will return the path where the file was saved
             download_path = await client.download_media(message, file=UPLOAD_DIR)
-
-            # Determine the filename to send to the user
-            response_filename = message.file.name if message.file and message.file.name else "downloaded_file"
-
-            # Return the file as a response
-            return FileResponse(download_path, media_type='application/octet-stream', filename=response_filename)
-
+            return FileResponse(download_path, media_type='application/octet-stream', filename=filename)
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to download file from Telegram: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to download from Telegram: {e}")
     finally:
-        # Crucially, clean up the downloaded file from the server's temp storage
         if download_path and os.path.exists(download_path):
             os.remove(download_path)