import os
import asyncio
from telethon.sync import TelegramClient
from telethon.sessions import StringSession
from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
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 and returns the full Telegram message object."""
    if not all([SESSION_STRING, API_ID, API_HASH]):
        raise ValueError("Server is not configured with Telegram credentials.")
    async with TelegramClient(StringSession(SESSION_STRING), API_ID, API_HASH) as client:
        try:
            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.")
            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(...)):
    """
    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())

        message = await upload_to_telegram(file_location)
        doc = message.document
        
        original_filename = next((attr.file_name for attr in doc.attributes if hasattr(attr, 'file_name')), file.filename)

        # Return details without the channel link, as requested.
        return {
            "message": "File uploaded successfully. Use the 'message_id' to download.",
            "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:
        if os.path.exists(file_location):
            os.remove(file_location)

@app.get("/v1/download/{message_id}")
async def download_file(message_id: int, background_tasks: BackgroundTasks):
    """
    Downloads a file from Telegram using its unique message_id from the channel.
    This method is stateless and resilient to server restarts.
    """
    if not all([SESSION_STRING, API_ID, API_HASH]):
        raise HTTPException(status_code=500, detail="Server is not configured for downloads.")

    try:
        async with TelegramClient(StringSession(SESSION_STRING), API_ID, API_HASH) as client:
            message = await client.get_messages(TARGET_CHANNEL, ids=message_id)
            
            if not message or not message.media:
                raise HTTPException(status_code=440, detail=f"A message with ID '{message_id}' was not found or contains no media.")

            original_filename = message.file.name if message.file and message.file.name else f"download_{message_id}"
            download_path = await client.download_media(message, file=UPLOAD_DIR)

            # **THE FIX IS HERE:**
            # The os.remove is now a background task. It will only run AFTER the
            # FileResponse has been fully sent to the client, preventing the error.
            if download_path:
                background_tasks.add_task(os.remove, download_path)
            
            return FileResponse(download_path, media_type='application/octet-stream', filename=original_filename)

    except Exception as e:
        # Catch potential errors like file not found in Telegram
        raise HTTPException(status_code=500, detail=f"Failed to download from Telegram: {e}")