Update main.py

main.py

@@ -5,71 +5,100 @@ from telethon.sessions import StringSession
 from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.responses import FileResponse
 import uvicorn
-import tempfile # Import the tempfile module
+import tempfile
 
 # --- Configuration from environment variables ---
-# Your Telegram session string
 SESSION_STRING = os.environ.get('SESSION_STRING')
-
-# You MUST get your own API ID and HASH from my.telegram.org
 API_ID = os.environ.get('API_ID')
 API_HASH = os.environ.get('API_HASH')
 
-# The Telegram channel you want to upload files to
+# The Telegram channel to use as file storage
 TARGET_CHANNEL = 'https://t.me/TestNAI01'
 
-# --- FIX: Use the system's temporary directory ---
-# This directory is guaranteed to be writable and avoids permission errors.
+# Use the system's temporary directory for transient files
 UPLOAD_DIR = tempfile.gettempdir()
 
 app = FastAPI()
 
 async def upload_to_telegram(file_path: str):
     """
-    Uploads a file to the specified Telegram channel using a session string.
+    Uploads a file to the Telegram channel and returns the 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.")
 
-    # Initialize the client
     async with TelegramClient(StringSession(SESSION_STRING), API_ID, API_HASH) as client:
         try:
-            await client.send_file(TARGET_CHANNEL, file_path)
+            # 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
+            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.")
+            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(...)):
     """
-    Receives a file, saves it to the temp directory, and uploads it to the Telegram channel.
+    Receives a file, uploads it to Telegram, and returns the Telegram message ID.
     """
-    # Create a secure path in the temporary directory
     file_location = os.path.join(UPLOAD_DIR, file.filename)
     try:
-        # Save the uploaded file temporarily
         with open(file_location, "wb+") as file_object:
             file_object.write(file.file.read())
 
-        # Upload the file to Telegram
-        await upload_to_telegram(file_location)
-        return {"message": f"File '{file.filename}' uploaded successfully to Telegram."}
+        # Upload the file and get the message object
+        message = await upload_to_telegram(file_location)
+
+        # 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
+
+        return {
+            "message": "File uploaded successfully to Telegram.",
+            "telegram_file_id": telegram_file_id,
+            "filename": original_filename,
+            "channel": TARGET_CHANNEL
+        }
     finally:
-        # Clean up the local file after uploading to save space
+        # Clean up the local temporary 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):
     """
-    Serves a file for download.
-    NOTE: This endpoint is for demonstration. It looks for files in the temp
-    directory, but files uploaded in previous requests will have been deleted.
-    A production system would need a proper way to retrieve files from Telegram.
+    Downloads a file from Telegram using its message ID.
     """
-    file_path = os.path.join(UPLOAD_DIR, filename)
-    if os.path.exists(file_path):
-        return FileResponse(file_path)
-    raise HTTPException(status_code=404, detail="File not found. Files are only stored temporarily during upload.")
+    if not all([SESSION_STRING, API_ID, API_HASH]):
+        raise HTTPException(status_code=500, detail="Server is not configured with Telegram credentials.")
+
+    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.")
+
+            # 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)
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to download file 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)
 
 if __name__ == "__main__":
     if not all([SESSION_STRING, API_ID, API_HASH]):