import gradio as gr
import tempfile, os
from pdf2image import convert_from_path
import pytesseract, pdfplumber, camelot
from PIL import Image, ImageOps
import ocrmypdf
import subprocess

def extract_text_from_pdf(file):
    extracted = []
    pdf_path = file.name
    
    # Create temporary paths for OCR'd PDF and text output
    temp_dir = tempfile.gettempdir()
    ocr_pdf_path = os.path.join(temp_dir, "ocr_searchable.pdf")
    output_txt_path = os.path.join(temp_dir, "extracted_text.txt")
    
    try:
        # Step 1: Use OCRmyPDF to create a searchable PDF
        print("Processing PDF with OCRmyPDF...")
        ocrmypdf.ocr(
            pdf_path, 
            ocr_pdf_path, 
            deskew=True,
            clean=True,
            force_ocr=False,  # Only OCR if needed
            skip_text=False,
            optimize=1
        )
        
        # Step 2: Extract text from the OCR'd searchable PDF using pdfplumber
        print("Extracting text from OCR'd PDF...")
        with pdfplumber.open(ocr_pdf_path) as pdf:
            for page_num, page in enumerate(pdf.pages):
                text = page.extract_text(layout=True)
                if text:
                    extracted.append(f"--- Page {page_num + 1} ---\n{text}")
                
                # Extract tables if any
                tables = page.extract_tables()
                for table_num, table in enumerate(tables):
                    if table:
                        table_text = f"TABLE {table_num + 1} (Page {page_num + 1}):\n"
                        table_text += "\n".join([", ".join([str(cell) if cell else "" for cell in row]) for row in table])
                        extracted.append(table_text)
        
        # Step 3: Try Camelot for additional table extraction
        try:
            tables = camelot.read_pdf(ocr_pdf_path, pages="all", flavor="lattice")
            for i, table in enumerate(tables):
                extracted.append(f"CAMELOT TABLE {i + 1}:\n{table.df.to_csv(index=False)}")
        except Exception as e:
            print(f"Camelot extraction failed: {e}")
        
        # Combine all extracted text
        combined_text = "\n\n".join(extracted).strip()
        
        # If still no text, fallback to direct OCR
        if len(combined_text) < 50:
            print("Fallback to direct OCR...")
            images = convert_from_path(pdf_path, dpi=300)
            ocr_text = []
            for i, img in enumerate(images):
                img = img.convert("L")
                img = ImageOps.invert(img)
                page_text = pytesseract.image_to_string(img, config="--psm 6")
                if page_text.strip():
                    ocr_text.append(f"--- Page {i + 1} (Direct OCR) ---\n{page_text}")
            combined_text = "\n\n".join(ocr_text)
        
        # Save the extracted text
        with open(output_txt_path, "w", encoding="utf-8") as f:
            f.write(combined_text)
        
        return combined_text, output_txt_path, ocr_pdf_path
        
    except Exception as e:
        error_msg = f"Error processing PDF: {str(e)}\n\nFalling back to original extraction methods..."
        print(error_msg)
        
        # Fallback to original method if OCRmyPDF fails
        try:
            with pdfplumber.open(pdf_path) as pdf:
                for page in pdf.pages:
                    text = page.extract_text(layout=True)
                    if text:
                        extracted.append(text)
                    tables = page.extract_tables()
                    for table in tables:
                        extracted.append("TABLE:\n" + "\n".join([", ".join(row) for row in table]))
        except Exception as e2:
            print("pdfplumber error:", e2)

        # OCR fallback if text is too short
        combined = "\n".join(extracted).strip()
        if len(combined) < 100:
            images = convert_from_path(pdf_path, dpi=300)
            for img in images:
                img = img.convert("L")
                img = ImageOps.invert(img)
                combined += pytesseract.image_to_string(img, config="--psm 6") + "\n"

        # Save fallback output
        with open(output_txt_path, "w", encoding="utf-8") as f:
            f.write(combined)

        return combined, output_txt_path, pdf_path  # Return original PDF if OCR failed

# Create Gradio interface
app = gr.Interface(
    fn=extract_text_from_pdf,
    inputs=gr.File(label="📤 Upload PDF", file_types=[".pdf"]),
    outputs=[
        gr.Textbox(label="📄 Extracted Text", lines=25, show_copy_button=True),
        gr.File(label="📥 Download Extracted Text (.txt)"),
        gr.File(label="📥 Download OCR'd Searchable PDF")
    ],
    title="Advanced PDF OCR Extractor with OCRmyPDF",
    description="Upload a PDF to get: 1) Extracted text displayed and downloadable as .txt, 2) OCR'd searchable PDF download. Uses OCRmyPDF for superior OCR quality.",
    allow_flagging="never",
)

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