main.py

import os
import asyncio
from telethon.sync import TelegramClient
from telethon.sessions import StringSession
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import StreamingResponse
import uvicorn
import tempfile

# --- Configuration from environment variables ---
SESSION_STRING = os.environ.get('SESSION_STRING')
API_ID = os.environ.get('API_ID')
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 uploads
UPLOAD_DIR = tempfile.gettempdir()

app = FastAPI()

@app.post("/v1/upload")
async def upload_file(file: UploadFile = File(...)):
    """
    Uploads a file to Telegram and returns its persistent message_id and other 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())

        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)

            # Simplified response with only the essentials
            return {
                "message": "File uploaded successfully. Use the 'message_id' to download.",
                "telegram_details": {
                    "message_id": message.id,
                    "filename": original_filename,
                    "size_bytes": doc.size,
                    "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):
    """
    Streams a file directly from Telegram to the client, ensuring a fast download start.
    """
    if not all([SESSION_STRING, API_ID, API_HASH]):
        raise HTTPException(status_code=500, detail="Server is not configured for downloads.")

    client = TelegramClient(StringSession(SESSION_STRING), API_ID, API_HASH)
    await client.connect()
    
    try:
        message = await client.get_messages(TARGET_CHANNEL, ids=message_id)
        if not message or not message.media:
            raise HTTPException(status_code=404, detail=f"Message not found or has no media.")
    except Exception as e:
        await client.disconnect()
        raise HTTPException(status_code=500, detail=f"Could not retrieve file details: {e}")

    original_filename = message.file.name if message.file and message.file.name else f"download_{message_id}"
    mime_type = message.file.mime_type if message.file else 'application/octet-stream'
    
    headers = {
        'Content-Disposition': f'attachment; filename="{original_filename}"'
    }

    async def file_iterator():
        """This async generator streams the file and ensures the client disconnects."""
        try:
            # Stream the file chunk by chunk
            async for chunk in client.iter_download(message):
                yield chunk
        finally:
            # Crucially, disconnect after the download is complete or fails
            await client.disconnect()

    return StreamingResponse(file_iterator(), media_type=mime_type, headers=headers)

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)