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 FileResponse
import uvicorn
import tempfile
from typing import Dict

# --- 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()

# --- 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 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, 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())

        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)

        # 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.",
            "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:
        if os.path.exists(file_location):
            os.remove(file_location)

@app.get("/v1/download/{filename}")
async def download_file(filename: str):
    """
    Downloads a file from Telegram by its filename using the stored mapping.
    """
    # 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:
            message = await client.get_messages(TARGET_CHANNEL, ids=message_id)
            if not message or not message.media:
                raise HTTPException(status_code=404, detail="File found in map, but message is missing or has no media in Telegram.")

            download_path = await client.download_media(message, file=UPLOAD_DIR)
            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 from Telegram: {e}")
    finally:
        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)