import os
import uuid
import io
import shutil
from PIL import Image
import pytesseract
import fitz  # PyMuPDF
import gradio as gr
from pptx import Presentation
from pptx.util import Inches, Pt
from docx import Document
from docx.shared import Cm

# -------------------------
# Configuration
# -------------------------
UPLOAD_DIR = "uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)
MAX_FILE_SIZE_MB = 500

# Explicit path to tesseract (HF Spaces Linux default)
pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract"

# -------------------------
# OCR Extraction
# -------------------------
def ocr_page_image(image: Image.Image):
    """
    Perform OCR on a PIL image and return list of text blocks.
    """
    data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
    text_blocks = []
    n_boxes = len(data['level'])
    for i in range(n_boxes):
        text = data['text'][i].strip()
        if text:
            # Capture bounding box
            block = {
                'text': text,
                'left': data['left'][i],
                'top': data['top'][i],
                'width': data['width'][i],
                'height': data['height'][i]
            }
            text_blocks.append(block)
    return text_blocks

# -------------------------
# PDF → Word
# -------------------------
def pdf_to_word_ocr(pdf_doc, output_path):
    word_doc = Document()
    for page in pdf_doc:
        pix = page.get_pixmap()
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        blocks = ocr_page_image(img)
        page_text = "\n".join([b['text'] for b in blocks])
        word_doc.add_paragraph(page_text)
        word_doc.add_page_break()
    word_doc.save(output_path)

# -------------------------
# PDF → PowerPoint
# -------------------------
def pdf_to_ppt_ocr(pdf_doc, output_path):
    prs = Presentation()
    prs.slide_width = Inches(13.333)
    prs.slide_height = Inches(7.5)

    for page in pdf_doc:
        pix = page.get_pixmap()
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        blocks = ocr_page_image(img)

        slide = prs.slides.add_slide(prs.slide_layouts[6])
        # Place text boxes for each OCR block
        for block in blocks:
            left = Inches(block['left'] / 100)  # scale roughly
            top = Inches(block['top'] / 100)
            width = Inches(block['width'] / 100)
            height = Inches(block['height'] / 100)
            textbox = slide.shapes.add_textbox(left, top, width, height)
            tf = textbox.text_frame
            p = tf.add_paragraph()
            p.text = block['text']
            p.font.size = Pt(14)
        # Optional: insert full page image as reference below text
        img_stream = io.BytesIO()
        img.save(img_stream, format="PNG")
        img_stream.seek(0)
        slide.shapes.add_picture(img_stream, 0, 0, width=prs.slide_width, height=prs.slide_height)
        img_stream.close()

    prs.save(output_path)

# -------------------------
# Main Conversion
# -------------------------
def convert_pdf(file_path, target_format):
    if file_path is None:
        raise gr.Error("Please upload a PDF file.")

    if not file_path.lower().endswith(".pdf"):
        raise gr.Error("Only PDF files are supported.")

    unique_id = str(uuid.uuid4())
    pdf_path = os.path.join(UPLOAD_DIR, f"{unique_id}.pdf")
    shutil.copy(file_path, pdf_path)

    pdf_doc = fitz.open(pdf_path)
    if len(pdf_doc) == 0:
        pdf_doc.close()
        os.remove(pdf_path)
        raise gr.Error("The uploaded PDF is empty.")

    output_file = None
    try:
        if target_format == "Word Document (.docx)":
            output_file = os.path.join(UPLOAD_DIR, f"{unique_id}.docx")
            pdf_to_word_ocr(pdf_doc, output_file)
        else:
            output_file = os.path.join(UPLOAD_DIR, f"{unique_id}.pptx")
            pdf_to_ppt_ocr(pdf_doc, output_file)
    finally:
        pdf_doc.close()
        os.remove(pdf_path)

    return output_file

# -------------------------
# Clear Uploads
# -------------------------
def clear_files():
    for f in os.listdir(UPLOAD_DIR):
        try:
            os.remove(os.path.join(UPLOAD_DIR, f))
        except:
            pass
    return None

# -------------------------
# Gradio UI
# -------------------------
with gr.Blocks(title="PDF Converter | OCR Editable") as demo:
    gr.Markdown("""
    # 📄 PDF Converter (OCR for Image-Based PDFs)
    Convert scanned or image-based PDFs into **editable Word or PowerPoint** files.
    """)

    with gr.Row():
        pdf_input = gr.File(
            label="Upload PDF File",
            file_types=[".pdf"]
        )
        format_choice = gr.Radio(
            choices=["Word Document (.docx)", "PowerPoint (.pptx)"],
            value="Word Document (.docx)",
            label="Select Output Format"
        )

    with gr.Row():
        convert_btn = gr.Button("Convert")
        clear_btn = gr.Button("Clear Uploaded Files")

    output_file = gr.File(label="Download Converted File")

    convert_btn.click(
        fn=convert_pdf,
        inputs=[pdf_input, format_choice],
        outputs=output_file
    )
    clear_btn.click(
        fn=clear_files,
        outputs=output_file
    )

# -------------------------
# Launch
# -------------------------
demo.launch(server_name="0.0.0.0", server_port=7860)