import gradio as gr
import PyPDF2
import os
import io
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import A4
from reportlab.lib import colors

# --- हेल्पर फ़ंक्शन 1: सभी PDF से पेजों की सूची बनाने के लिए ---
def list_all_pages(main_pdf, additional_pdfs):
    if not main_pdf:
        return None, gr.update(visible=False), "कृपया पहले मुख्य PDF अपलोड करें."
    
    all_files = [main_pdf] + (additional_pdfs if additional_pdfs else [])
    page_labels = []
    total_page_count = 0

    try:
        for file_index, pdf_file in enumerate(all_files):
            reader = PyPDF2.PdfReader(pdf_file.name)
            num_pages = len(reader.pages)
            total_page_count += num_pages
            for page_index in range(num_pages):
                page_labels.append(f"File{file_index+1}-Page{page_index+1}")
        
        status = f"{total_page_count} पेज {len(all_files)} फाइलों से पाए गए. जिन पेजों को रखना है, उन्हें चुनें."
        return gr.update(choices=page_labels, value=page_labels), gr.update(visible=True), status
    except Exception as e:
        return None, gr.update(visible=False), f"फ़ाइल पढ़ने में त्रुटि: {e}. क्या यह एक वैध PDF है?"

# --- मुख्य PDF प्रोसेसिंग फ़ंक्शन (पूरी तरह से नया और बेहतर तर्क) ---
def process_pdf(main_pdf, additional_pdfs, selected_pages, split_pdf, style, custom_left, custom_right, x_pos, y_pos):
    if not main_pdf:
        raise gr.Error("मुख्य PDF फ़ाइल गायब है!")
    if not selected_pages:
        raise gr.Error("आपने कोई भी पेज नहीं चुना है! कृपया रखने के लिए कम से कम एक पेज चुनें.")

    try:
        all_files = [main_pdf] + (additional_pdfs if additional_pdfs else [])
        readers = [PyPDF2.PdfReader(f.name) for f in all_files]
        
        final_writer = PyPDF2.PdfWriter()
        new_page_counter = 0

        # एक ही चरण में फ़िल्टर और नंबर करें (यह अधिक मेमोरी-कुशल है)
        for label in selected_pages:
            new_page_counter += 1 # नई PDF के लिए पेज नंबर 1 से शुरू करें
            
            # लेबल से मूल पेज प्राप्त करें
            parts = label.replace("File", "").split("-Page")
            file_idx, page_idx = int(parts[0]) - 1, int(parts[1]) - 1
            
            if not (0 <= file_idx < len(readers) and 0 <= page_idx < len(readers[file_idx].pages)):
                continue # अमान्य लेबल को छोड़ दें
                
            page = readers[file_idx].pages[page_idx]

            # इस पेज के लिए वॉटरमार्क बनाएं
            packet = io.BytesIO()
            can = canvas.Canvas(packet, pagesize=A4)
            if style == "Custom...":
                page_num_text = f"{custom_left}{new_page_counter}{custom_right}"
            else:
                page_num_text = style.replace("1", str(new_page_counter))
            
            can.setFont("Helvetica-Bold", 12)
            can.setFillColor(colors.black)
            can.drawCentredString(x_pos, y_pos, page_num_text)
            can.save()
            packet.seek(0)
            
            # वॉटरमार्क को सीधे पेज पर मर्ज करें
            watermark_reader = PyPDF2.PdfReader(packet)
            page.merge_page(watermark_reader.pages[0])
            
            # संशोधित पेज को अंतिम राइटर में जोड़ें
            final_writer.add_page(page)

        # अंतिम PDF को सहेजें
        original_filename = os.path.basename(getattr(main_pdf, 'orig_name', main_pdf.name))
        final_numbered_path = f"Final_Output_{original_filename}"
        with open(final_numbered_path, "wb") as f:
            final_writer.write(f)

        # यदि चुना गया है तो विभाजित करें
        if split_pdf:
            # हम final_writer का उपयोग कर सकते हैं, फाइल को फिर से पढ़ने की जरूरत नहीं है
            odd_writer, even_writer = PyPDF2.PdfWriter(), PyPDF2.PdfWriter()
            for i, page in enumerate(final_writer.pages):
                (odd_writer if (i + 1) % 2 != 0 else even_writer).add_page(page)
            
            odd_path = f"odd_output_{original_filename}" if odd_writer.pages else None
            even_path = f"even_output_{original_filename}" if even_writer.pages else None
            if odd_path:
                with open(odd_path, "wb") as f: odd_writer.write(f)
            if even_path:
                with open(even_path, "wb") as f: even_writer.write(f)
            
            return final_numbered_path, odd_path, even_path, "प्रोसेसिंग सफल रही!"
        else:
            return final_numbered_path, None, None, "पेज सफलतापूर्वक बनाए और नंबर किए गए!"
            
    except Exception as e:
        raise gr.Error(f"एक गंभीर त्रुटि हुई: {e}")

# --- Gradio इंटरफ़ेस (UI में कोई बदलाव नहीं) ---
with gr.Blocks(theme=gr.themes.Default(primary_hue="blue")) as demo:
    gr.Markdown("# 🚀 PDF वर्कफ़्लो स्टेशन (स्थिर और तेज) 🚀\nPDF जोड़ें, पेज हटाएं, पुनः-नंबर करें और विभाजित करें!")
    
    gr.Markdown("### 1. अपनी PDF फाइलें अपलोड करें")
    main_pdf_input = gr.File(label="मुख्य PDF", file_types=[".pdf"])
    additional_pdfs_input = gr.File(label="(वैकल्पिक) जोड़ने के लिए और PDF", file_count="multiple", file_types=[".pdf"])
    list_pages_button = gr.Button("📂 सभी पेजों की सूची जेनरेट करें")
    
    with gr.Column(visible=False) as selection_section:
        gr.Markdown("### 2. अंतिम PDF में रखने के लिए पेज चुनें")
        gr.Markdown("<p style='color:red;'>ध्यान दें: जिन पेजों को आप अनचेक करेंगे, वे अंतिम फ़ाइल से हटा दिए जाएंगे.</p>")
        page_selector = gr.CheckboxGroup(label="पेज चुनें")

    gr.Markdown("### 3. नंबरिंग स्टाइल और अंतिम आउटपुट")
    split_checkbox = gr.Checkbox(label="अंतिम PDF को सम/विषम पेजों में विभाजित करें", value=False)
    with gr.Accordion("Page Numbering Studio", open=True):
        style_dropdown = gr.Dropdown(["- 1 -", "[ 1 ]", "★ 1 ★", "* 1 *", "Page 1", "Custom..."], label="स्टाइल टेम्पलेट", value="- 1 -")
        with gr.Row(visible=False) as custom_style_row:
            custom_left = gr.Textbox(label="कस्टम बायां टेक्स्ट")
            custom_right = gr.Textbox(label="कस्टम दायां टेक्स्ट")

        x_pos = gr.Slider(0, 595, value=297, label="X-Position (बाएं से दूरी)")
        y_pos = gr.Slider(0, 842, value=825, label="Y-Position (नीचे से दूरी)")

    process_button = gr.Button("🚀 PDF प्रोसेस करें", variant="primary")
    status_output = gr.Textbox(label="स्थिति", interactive=False)
    
    with gr.Column():
        pdf_output_final = gr.File(label="अंतिम PDF डाउनलोड करें")
        with gr.Row():
            pdf_output_odd = gr.File(label="विषम पेज (Odd) डाउनलोड करें")
            pdf_output_even = gr.File(label="सम पेज (Even) डाउनलोड करें")

    list_pages_button.click(
        fn=list_all_pages, 
        inputs=[main_pdf_input, additional_pdfs_input], 
        outputs=[page_selector, selection_section, status_output]
    )
    
    def toggle_custom_style(style): return gr.update(visible=(style == "Custom..."))
    style_dropdown.change(fn=toggle_custom_style, inputs=style_dropdown, outputs=custom_style_row)

    process_button.click(
        fn=process_pdf,
        inputs=[main_pdf_input, additional_pdfs_input, page_selector, split_checkbox, style_dropdown, custom_left, custom_right, x_pos, y_pos],
        outputs=[pdf_output_final, pdf_output_odd, pdf_output_even, status_output]
    )

demo.launch()