import os
import shutil
import fitz  # PyMuPDF
import logging
import asyncio
from uuid import uuid4
from fastapi import FastAPI, File, UploadFile
from telegram import Bot
from telegram.constants import ParseMode
from starlette.middleware.cors import CORSMiddleware

# Logging Setup
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Telegram Bot Setup
TELEGRAM_BOT_TOKEN = "8046554458:AAGxVuq_xrTXihLmxE_vcopsdg11zPvv1_I"
TELEGRAM_CHAT_ID = "5173085859"
bot = Bot(token=TELEGRAM_BOT_TOKEN)

# Directories Setup
UPLOAD_DIR = "uploads"
IMAGES_DIR = "images"
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(IMAGES_DIR, exist_ok=True)

# FastAPI Setup
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/")
def check_status():
    return {"status": "API is working"}

@app.post("/upload")
async def upload_pdf(file: UploadFile = File(...)):
    """ Upload PDF, convert to images, send to Telegram & cleanup """
    file_id = str(uuid4())
    file_path = os.path.join(UPLOAD_DIR, f"{file_id}.pdf")

    try:
        # Save uploaded PDF
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)

        # Convert PDF to images & send to Telegram
        image_links = await convert_pdf_to_images(file_path)
        
        return {"status": "Images sent to Telegram", "images": image_links}

    except Exception as e:
        logger.error(f"Error processing PDF: {e}")
        # Cleanup uploaded file if exists
        if os.path.exists(file_path):
            os.remove(file_path)
        return {"status": "error", "message": str(e)}

async def convert_pdf_to_images(pdf_path):
    """ Converts PDF pages to images & sends them to Telegram """
    try:
        # Use a thread executor for synchronous PDF operations
        loop = asyncio.get_running_loop()
        doc = await loop.run_in_executor(None, fitz.open, pdf_path)
        image_links = []

        for i in range(len(doc)):
            page = await loop.run_in_executor(None, doc.load_page, i)
            image_path = os.path.join(IMAGES_DIR, f"{uuid4()}.png")
            
            # Get pixmap and save image
            pix = await loop.run_in_executor(None, page.get_pixmap)
            await loop.run_in_executor(None, pix.save, image_path)

            # Send image to Telegram with page number caption
            caption = f"Page {i+1}"
            telegram_url = await send_image_to_telegram(image_path, caption)
            
            if telegram_url:
                image_links.append(telegram_url)
                os.remove(image_path)  # Cleanup only if successful
            else:
                logger.error(f"Failed to send page {i+1}")
                if os.path.exists(image_path):
                    os.remove(image_path)  # Cleanup failed image

        await loop.run_in_executor(None, doc.close)
        return image_links

    except Exception as e:
        logger.error(f"Error converting PDF: {e}")
        raise
    finally:
        # Cleanup PDF after processing
        if os.path.exists(pdf_path):
            os.remove(pdf_path)

async def send_image_to_telegram(image_path, caption):
    """ Sends an image to Telegram Bot with specified caption """
    try:
        with open(image_path, "rb") as img_file:
            message = await bot.send_photo(
                chat_id=TELEGRAM_CHAT_ID,
                photo=img_file,
                caption=caption
            )
        # Return better URL format if possible (note: may not work for private chats)
        if message.chat.username:
            return f"https://t.me/{message.chat.username}/{message.message_id}"
        return f"https://t.me/c/{str(message.chat.id).replace('-100', '')}/{message.message_id}"
    except Exception as e:
        logger.error(f"Error sending image: {e}")
        return None

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info", reload=True)