import gradio as gr
import fitz  # PyMuPDF
import arabic_reshaper
from bidi.algorithm import get_display
import pytesseract
from PIL import Image
import io
import numpy as np
import tempfile
import os

def extract_text_from_pdf(pdf_file):
    try:
        # ایجاد فایل موقت برای PDF
        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
            if hasattr(pdf_file, 'read'):
                # اگر فایل قابل خواندن است
                tmp_file.write(pdf_file.read())
            else:
                # اگر از Gradio File component است
                tmp_file.write(pdf_file)
            
            tmp_path = tmp_file.name
        
        # باز کردن فایل PDF از مسیر
        doc = fitz.open(tmp_path)
        full_text = ""
        has_ocr_processed = False
        
        # استخراج متن از تمام صفحات
        for page_num in range(len(doc)):
            page = doc.load_page(page_num)
            
            # ابتدا سعی می‌کنیم متن را مستقیماً استخراج کنیم
            text = page.get_text()
            
            # اگر متن کمی استخراج شده (صفحه احتمالاً تصویری است)
            if len(text.strip()) < 100:
                # تبدیل صفحه به تصویر با وضوح بالا
                mat = fitz.Matrix(300/72, 300/72)  # وضوح 300 DPI
                pix = page.get_pixmap(matrix=mat)
                img_data = pix.tobytes("png")
                
                # استفاده از OCR برای استخراج متن از تصویر
                image = Image.open(io.BytesIO(img_data))
                
                # تنظیمات OCR برای زبان‌های مختلف
                custom_config = r'--oem 3 --psm 6'
                ocr_text = pytesseract.image_to_string(
                    image, 
                    lang='fas+ara+eng',
                    config=custom_config
                )
                text = ocr_text
                has_ocr_processed = True
            
            full_text += f"--- صفحه {page_num + 1} ---\n"
            full_text += text + "\n\n"
        
        doc.close()
        
        # حذف فایل موقت
        os.unlink(tmp_path)
        
        # پردازش متن برای زبان‌های راست‌به‌چپ
        try:
            reshaped_text = arabic_reshaper.reshape(full_text)
            bidi_text = get_display(reshaped_text)
            
            if has_ocr_processed:
                bidi_text = "[⚠️ برخی صفحات با OCR پردازش شدند]\n\n" + bidi_text
            
            return bidi_text
        except Exception as proc_error:
            print(f"خطا در پردازش متن: {proc_error}")
            return full_text
            
    except Exception as e:
        # حذف فایل موقت در صورت خطا
        try:
            if 'tmp_path' in locals():
                os.unlink(tmp_path)
        except:
            pass
        return f"خطا در پردازش فایل: {str(e)}"

# ایجاد رابط Gradio
with gr.Blocks(title="PDF Text Extractor with OCR") as demo:
    gr.Markdown("# 📄 استخراج متن از فایل PDF")
    gr.Markdown("با قابلیت پردازش OCR برای PDFهای تصویری")
    
    with gr.Row():
        pdf_input = gr.File(
            label="فایل PDF را انتخاب کنید", 
            file_types=[".pdf"],
            type="bytes"  # استفاده از bytes به جای file path
        )
    
    extract_btn = gr.Button("🔄 استخراج متن")
    
    with gr.Row():
        text_output = gr.Textbox(
            label="متن استخراج شده", 
            lines=20, 
            interactive=False,
            show_copy_button=True
        )
    
    gr.Markdown("""
    **⚠️ توجه:**  
    - برای PDFهای متنی، متن مستقیماً استخراج می‌شود  
    - برای PDFهای تصویری، از OCR استفاده می‌شود  
    - پشتیبانی از زبان‌های فارسی، عربی و انگلیسی  
    - پردازش ممکن است برای فایل‌های بزرگ کمی زمان‌بر باشد  
    """)
    
    extract_btn.click(
        fn=extract_text_from_pdf,
        inputs=pdf_input,
        outputs=text_output
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)