import os
import uuid
import time
import asyncio
import fitz  # PyMuPDF
import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI, Request, Response
from fastapi.responses import FileResponse
import uvicorn

from telegram import Update, ReplyKeyboardMarkup, KeyboardButton
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes

# --- تنظیمات لاگ و متغیرها ---
logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)
logger = logging.getLogger(__name__)

TOKEN = os.environ.get("TOKEN")
DOWNLOAD_DIR = "downloads"
SPACE_URL = "https://barnamemaker-pdfmeger.hf.space"
MERGED_FILE_TTL = 3600  # ۱ ساعت

os.makedirs(DOWNLOAD_DIR, exist_ok=True)
user_files: dict[int, list[str]] = {}
user_locks: dict[int, asyncio.Lock] = {}

def get_user_lock(user_id: int) -> asyncio.Lock:
    if user_id not in user_locks:
        user_locks[user_id] = asyncio.Lock()
    return user_locks[user_id]

# --- تسک پاکسازی در پس‌زمینه ---
async def cleanup_task():
    while True:
        await asyncio.sleep(600)
        now = time.time()
        for f in os.listdir(DOWNLOAD_DIR):
            path = os.path.join(DOWNLOAD_DIR, f)
            try:
                if os.path.isfile(path) and (now - os.path.getmtime(path) > MERGED_FILE_TTL):
                    os.remove(path)
                    logger.info(f"Cleaned up old file: {f}")
            except Exception as e:
                pass

# --- منوی دکمه‌ها ---
def get_menu():
    keyboard = [
        [KeyboardButton("📄 ادغام فایل‌ها")],
        [KeyboardButton("🗑 پاک کردن صف"), KeyboardButton("❓ راهنما")]
    ]
    return ReplyKeyboardMarkup(keyboard, resize_keyboard=True, one_time_keyboard=False)

HELP_TEXT = (
    "📖 راهنمای استفاده:\n\n"
    "۱. فایل‌های PDF خود را به هر تعدادی که می‌خواهید یکی‌یکی بفرستید.\n"
    "۲. روی «📄 ادغام فایل‌ها» کلیک کنید.\n"
    "۳. لینک دانلود فایل ادغام‌شده را دریافت کنید.\n\n"
    "برای پاک کردن صف: «🗑 پاک کردن صف»\n"
)

# --- منطق ربات (Handlers) ---
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    logger.info(f"User {user_id} started the bot.")
    await update.message.reply_text(
        "سلام! 👋\nفایل‌های PDF خود را بفرستید و سپس «📄 ادغام فایل‌ها» را بزنید.",
        reply_markup=get_menu()
    )

async def handle_document(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    doc = update.message.document

    if doc.mime_type != "application/pdf":
        await update.message.reply_text("⚠️ فقط فایل PDF قبول می‌شود.", reply_markup=get_menu())
        return

    lock = get_user_lock(user_id)
    async with lock:
        if user_id not in user_files:
            user_files[user_id] = []

        logger.info(f"User {user_id} sent a document: {doc.file_name}")
        file = await doc.get_file()
        file_id = f"{user_id}_{uuid.uuid4().hex[:8]}.pdf"
        file_path = os.path.join(DOWNLOAD_DIR, file_id)

        try:
            await file.download_to_drive(file_path)
        except Exception as e:
            logger.error(f"Download error: {e}")
            await update.message.reply_text("❌ خطا در دریافت فایل.", reply_markup=get_menu())
            return

        try:
            with fitz.open(file_path) as test_doc:
                if test_doc.page_count == 0:
                    raise ValueError("Empty PDF")
        except Exception:
            os.remove(file_path)
            await update.message.reply_text("⚠️ فایل PDF معتبر نیست.", reply_markup=get_menu())
            return

        user_files[user_id].append(file_path)
        count = len(user_files[user_id])

    await update.message.reply_text(f"✅ فایل {count} دریافت شد: {doc.file_name}", reply_markup=get_menu())

async def merge_pdfs(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    lock = get_user_lock(user_id)

    async with lock:
        files = list(user_files.get(user_id, []))

    if not files:
        await update.message.reply_text("⚠️ صف شما خالی است!", reply_markup=get_menu())
        return

    if len(files) == 1:
        await update.message.reply_text("⚠️ حداقل ۲ فایل برای ادغام لازم است.", reply_markup=get_menu())
        return

    msg = await update.message.reply_text(f"⏳ در حال ادغام {len(files)} فایل...")
    output_filename = f"merged_{user_id}_{uuid.uuid4().hex[:8]}.pdf"
    output_path = os.path.join(DOWNLOAD_DIR, output_filename)

    try:
        result = fitz.open()
        for pdf_path in files:
            if not os.path.exists(pdf_path): continue
            with fitz.open(pdf_path) as doc:
                result.insert_pdf(doc)

        result.save(output_path)
        result.close()

        async with lock:
            for f in user_files.get(user_id, []):
                if os.path.exists(f): os.remove(f)
            user_files[user_id] = []

        link = f"{SPACE_URL}/download/{output_filename}"
        await msg.edit_text(f"🎉 ادغام موفق!\n📥 لینک دانلود:\n{link}")

    except Exception as e:
        logger.error(f"Error merging: {e}")
        await msg.edit_text("❌ خطایی در پردازش رخ داد.")

async def clear_queue(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    lock = get_user_lock(user_id)
    async with lock:
        files = list(user_files.get(user_id, []))
        user_files[user_id] = []
    
    for f in files:
        if os.path.exists(f): os.remove(f)
    await update.message.reply_text("🗑 صف پاک شد.", reply_markup=get_menu())

async def handle_text(update: Update, context: ContextTypes.DEFAULT_TYPE):
    text = update.message.text
    if text == "📄 ادغام فایل‌ها": await merge_pdfs(update, context)
    elif text == "🗑 پاک کردن صف": await clear_queue(update, context)
    elif text == "❓ راهنما": await update.message.reply_text(HELP_TEXT, reply_markup=get_menu())

# --- تنظیمات اتصال اصلی (Webhook) ---
ptb_app = Application.builder().token(TOKEN).connect_timeout(40.0).read_timeout(40.0).build()
ptb_app.add_handler(CommandHandler("start", start))
ptb_app.add_handler(MessageHandler(filters.Document.PDF, handle_document))
ptb_app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_text))

@asynccontextmanager
async def lifespan(app: FastAPI):
    # این تابع در پس‌زمینه اجرا می‌شود تا سرور وب منتظر نماند
    async def setup_telegram():
        await asyncio.sleep(3)  # صبر می‌کنیم تا سرور وب کاملاً بالا بیاید
        try:
            logger.info("⏳ در حال راه‌اندازی ربات در پس‌زمینه...")
            await ptb_app.initialize()
            await ptb_app.bot.set_webhook(
                url=f"{SPACE_URL}/webhook", 
                drop_pending_updates=True,
                connect_timeout=60.0,
                read_timeout=60.0
            )
            await ptb_app.start()
            logger.info("✅ Webhook با موفقیت ثبت شد و ربات آماده است!")
            asyncio.create_task(cleanup_task())
        except Exception as e:
            logger.error(f"❌ خطا در تنظیم Webhook: {e}")

    # استارت تسک در پس‌زمینه
    setup_task = asyncio.create_task(setup_telegram())
    
    yield
    
    # خاموش کردن امن
    try:
        setup_task.cancel()
        await ptb_app.stop()
        await ptb_app.shutdown()
    except Exception:
        pass

app_web = FastAPI(lifespan=lifespan)

@app_web.post("/webhook")
async def telegram_webhook(request: Request):
    req_json = await request.json()
    update = Update.de_json(req_json, ptb_app.bot)
    await ptb_app.process_update(update)
    return Response(status_code=200)

@app_web.get("/")
async def root():
    return {"message": "Bot Webhook is Active!"}

@app_web.get("/download/{filename}")
async def get_file(filename: str):
    if ".." in filename or "/" in filename: return {"error": "Invalid"}
    file_path = os.path.join(DOWNLOAD_DIR, filename)
    if os.path.exists(file_path):
        return FileResponse(file_path, media_type="application/pdf")
    return {"error": "Not found"}

if __name__ == "__main__":
    if not TOKEN:
        raise RuntimeError("TOKEN is missing!")
    uvicorn.run("app:app_web", host="0.0.0.0", port=7860)