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 FileResponse
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 files
UPLOAD_DIR = tempfile.gettempdir()

app = FastAPI()

async def upload_to_telegram(file_path: str):
    """
    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.")

    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
            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, uploads it to Telegram, and returns the Telegram message ID.
    """
    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)

        # 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 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):
    """
    Downloads a file from Telegram using its message ID.
    """
    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]):
        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)