import gradio as gr
from pypdf import PdfWriter, PdfReader
import tempfile
import os


def merge_pdfs(pdf_files, add_bookmarks: bool):
    """Merge multiple PDF files into one."""
    if not pdf_files or len(pdf_files) < 2:
        raise gr.Error("Please upload at least 2 PDF files to merge.")

    writer = PdfWriter()
    total_pages = 0

    for pdf_path in pdf_files:
        try:
            reader = PdfReader(pdf_path)
            num_pages = len(reader.pages)

            if add_bookmarks:
                writer.add_outline_item(
                    title=os.path.basename(pdf_path),
                    page_number=total_pages,
                )

            for page in reader.pages:
                writer.add_page(page)

            total_pages += num_pages
        except Exception as e:
            raise gr.Error(f"Error reading '{os.path.basename(pdf_path)}': {e}")

    output_file = tempfile.NamedTemporaryFile(
        delete=False, suffix=".pdf", prefix="merged_"
    )
    with open(output_file.name, "wb") as f:
        writer.write(f)

    file_count = len(pdf_files)
    summary = (
        f"✅ Successfully merged **{file_count} PDFs** into **{total_pages} pages**."
    )
    return output_file.name, summary


def get_pdf_info(pdf_files):
    """Return a summary table of uploaded PDFs."""
    if not pdf_files:
        return ""

    rows = []
    for i, path in enumerate(pdf_files, 1):
        try:
            reader = PdfReader(path)
            pages = len(reader.pages)
            meta = reader.metadata
            title = (meta.title or "—") if meta else "—"
            size_kb = round(os.path.getsize(path) / 1024, 1)
            rows.append(f"| {i} | {os.path.basename(path)} | {pages} | {size_kb} KB | {title} |")
        except Exception:
            rows.append(f"| {i} | {os.path.basename(path)} | ⚠️ Error | — | — |")

    header = "| # | File | Pages | Size | Title |\n|---|------|-------|------|-------|"
    return header + "\n" + "\n".join(rows)


# ── UI ────────────────────────────────────────────────────────────────────────

with gr.Blocks(
    title="PDF Merger",
    theme=gr.themes.Soft(primary_hue="violet"),
) as demo:

    gr.Markdown(
        """
        # 📄 PDF Merger
        Upload two or more PDF files, reorder them if needed, and download the merged result.
        """
    )

    with gr.Row():
        with gr.Column(scale=2):
            pdf_input = gr.File(
                label="Upload PDFs (drag to reorder)",
                file_count="multiple",
                file_types=[".pdf"],
                type="filepath",
            )
            add_bookmarks = gr.Checkbox(
                label="📑 Add bookmarks for each file",
                value=True,
            )
            merge_btn = gr.Button("🔀 Merge PDFs", variant="primary", size="lg")

        with gr.Column(scale=1):
            info_md = gr.Markdown(label="File info")

    with gr.Row():
        output_file = gr.File(label="⬇️ Download merged PDF", interactive=False)
        result_msg = gr.Markdown()

    # Update info table whenever files change
    pdf_input.change(fn=get_pdf_info, inputs=pdf_input, outputs=info_md)

    merge_btn.click(
        fn=merge_pdfs,
        inputs=[pdf_input, add_bookmarks],
        outputs=[output_file, result_msg],
    )

    gr.Markdown(
        """
        ---
        *Built with [Gradio](https://gradio.app) & [pypdf](https://pypdf.readthedocs.io)*
        """
    )

if __name__ == "__main__":
  demo.launch()