import os
import io
import traceback
import logging
import torch
import gradio as gr
import fitz  # PyMuPDF
from PIL import Image
import spaces  # للـ ZeroGPU

# إعداد نظام التسجيل
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# تعطيل مسار الفيديو داخل Transformers
os.environ["TRANSFORMERS_NO_TORCHVISION"] = "1"

# أسماء النماذج
BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen2-VL-2B-Instruct")
HOROOF_ADAPTER = os.environ.get("HOROOF_MODEL", "NaserNajeh/Horoof")

# متغيرات النموذج العامة
_model = None
_tokenizer = None
_img_proc = None
_model_loaded = False

def check_gpu_availability():
    """التحقق من توفر GPU وطباعة معلومات النظام"""
    if not torch.cuda.is_available():
        raise AssertionError("هذه النسخة تتطلب GPU (CUDA) مفعّل على الـSpace.")
    
    device_name = torch.cuda.get_device_name(0)
    memory_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
    
    logger.info(f"GPU متاح: {device_name}")
    logger.info(f"ذاكرة GPU: {memory_gb:.1f} GB")
    
    return device_name, memory_gb

def clear_gpu_cache():
    """تنظيف ذاكرة GPU"""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

def load_model_merged():
    """
    تحميل النموذج مع تحسينات إضافية وإدارة أفضل للذاكرة
    """
    global _model, _tokenizer, _img_proc, _model_loaded
    
    if _model_loaded and _model is not None:
        logger.info("النموذج محمّل بالفعل")
        return
    
    try:
        # التحقق من GPU
        device_name, memory_gb = check_gpu_availability()
        
        # استيراد المكتبات
        from transformers import (
            Qwen2VLForConditionalGeneration,
            AutoTokenizer,
            Qwen2VLImageProcessor
        )
        from peft import PeftModel
        
        logger.info("بدء تحميل النموذج...")
        
        # تحميل Tokenizer و ImageProcessor
        logger.info("تحميل المعالجات...")
        _tokenizer = AutoTokenizer.from_pretrained(
            BASE_MODEL, 
            trust_remote_code=False,
            use_fast=True  # تسريع التوكين
        )
        _img_proc = Qwen2VLImageProcessor.from_pretrained(
            BASE_MODEL, 
            trust_remote_code=False
        )
        
        # تحميل النموذج الأساسي مع تحسينات الذاكرة
        logger.info(f"تحميل النموذج الأساسي: {BASE_MODEL}")
        
        # تحسين تحميل النموذج حسب حجم الذاكرة
        if memory_gb >= 40:  # H100 أو A100
            model_kwargs = {
                "torch_dtype": torch.float16,
                "device_map": "auto",
                "low_cpu_mem_usage": True
            }
        else:  # GPUs أصغر
            model_kwargs = {
                "torch_dtype": torch.float16,
                "low_cpu_mem_usage": True
            }
        
        base = Qwen2VLForConditionalGeneration.from_pretrained(
            BASE_MODEL,
            **model_kwargs
        )
        
        # نقل إلى GPU إذا لم يتم تلقائياً
        if not hasattr(base, 'device') or str(base.device) == 'cpu':
            base = base.to("cuda")
        
        # تحميل ودمج LoRA
        logger.info(f"تحميل LoRA adapter: {HOROOF_ADAPTER}")
        peft_model = PeftModel.from_pretrained(base, HOROOF_ADAPTER)
        
        logger.info("دمج LoRA مع النموذج الأساسي...")
        _model = peft_model.merge_and_unload()
        
        # التأكد من وجود النموذج على GPU
        if not str(_model.device).startswith('cuda'):
            _model = _model.to("cuda")
        
        # تحسين للاستنتاج
        _model.eval()
        
        # تمكين optimizations
        if hasattr(_model, 'half'):
            _model = _model.half()
        
        _model_loaded = True
        logger.info("تم تحميل النموذج بنجاح!")
        
        # طباعة معلومات الذاكرة
        if torch.cuda.is_available():
            allocated = torch.cuda.memory_allocated() / 1024**3
            logger.info(f"ذاكرة GPU المستخدمة: {allocated:.2f} GB")
        
    except Exception as e:
        logger.error(f"خطأ في تحميل النموذج: {str(e)}")
        logger.error(traceback.format_exc())
        raise RuntimeError(f"تعذّر تحميل النموذج: {e}")

def pdf_to_images(pdf_bytes: bytes, dpi: int = 220, max_pages: int = 0):
    """
    تحويل PDF إلى صور PIL مع تحسينات
    """
    try:
        pages_imgs = []
        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
        total = doc.page_count
        
        logger.info(f"عدد صفحات PDF: {total}")
        
        # تحديد عدد الصفحات للمعالجة
        n_pages = total if (not max_pages or max_pages <= 0) else min(max_pages, total)
        
        for i in range(n_pages):
            try:
                page = doc.load_page(i)
                
                # تحسين جودة التحويل
                mat = fitz.Matrix(dpi/72, dpi/72)
                pix = page.get_pixmap(matrix=mat, alpha=False)
                
                # تحويل لـ PIL Image
                img_data = pix.samples
                img = Image.frombytes("RGB", [pix.width, pix.height], img_data)
                
                # تحسين حجم الصورة إذا كانت كبيرة جداً
                max_dimension = 2048
                if max(img.size) > max_dimension:
                    img.thumbnail((max_dimension, max_dimension), Image.Resampling.LANCZOS)
                    logger.info(f"تم تصغير الصفحة {i+1} إلى {img.size}")
                
                pages_imgs.append((i + 1, img))
                logger.info(f"تم تحويل الصفحة {i+1}/{n_pages}")
                
            except Exception as e:
                logger.error(f"خطأ في تحويل الصفحة {i+1}: {str(e)}")
                continue
        
        doc.close()
        logger.info(f"تم تحويل {len(pages_imgs)} صفحة بنجاح")
        return pages_imgs
        
    except Exception as e:
        logger.error(f"خطأ في فتح PDF: {str(e)}")
        raise

def ocr_page_gpu(pil_img: Image.Image, max_new_tokens: int = 1200) -> str:
    """
    OCR لصفحة واحدة مع تحسينات الأداء
    """
    try:
        # التأكد من تحميل النموذج
        load_model_merged()
        
        # تنظيف الذاكرة قبل المعالجة
        clear_gpu_cache()
        
        # رسالة المحادثة
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": pil_img},
                    {"type": "text", "text": "اقرأ النص العربي في الصورة كما هو دون أي تعديل أو تفسير."},
                ],
            }
        ]

        # توكين النص
        tok = _tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_tensors="pt"
        )

        # معالجة الصورة
        vis = _img_proc(images=[pil_img], return_tensors="pt")

        # تجهيز المدخلات
        inputs = {"input_ids": tok}
        inputs.update(vis)
        
        # نقل إلى GPU
        for k, v in inputs.items():
            if hasattr(v, "to"):
                inputs[k] = v.to("cuda")

        # التوليد مع تحسينات
        with torch.inference_mode():
            output_ids = _model.generate(
                **inputs, 
                max_new_tokens=max_new_tokens,
                do_sample=False,  # للحصول على نتائج مستقرة
                temperature=0.1,
                pad_token_id=_tokenizer.eos_token_id,
                use_cache=True
            )
        
        # فك التشفير
        generated_part = output_ids[0][len(inputs["input_ids"][0]):]
        result = _tokenizer.decode(generated_part, skip_special_tokens=True)
        
        # تنظيف النتيجة
        result = result.strip()
        
        # تنظيف الذاكرة بعد المعالجة
        clear_gpu_cache()
        
        return result if result else "لم يتم استخراج أي نص"
        
    except Exception as e:
        logger.error(f"خطأ في OCR: {str(e)}")
        clear_gpu_cache()
        return f"خطأ في معالجة الصورة: {str(e)}"

@spaces.GPU  # ضروري لـ ZeroGPU
def ocr_pdf(pdf_file, dpi, limit_pages):
    """الدالة الرئيسية مع تحسينات إضافية"""
    
    if pdf_file is None:
        return "❌ لم يتم رفع ملف PDF."
    
    try:
        logger.info("بدء معالجة PDF...")
        
        # قراءة البيانات
        if hasattr(pdf_file, 'read'):
            pdf_bytes = pdf_file.read()
        else:
            pdf_bytes = pdf_file
        
        if not pdf_bytes:
            return "❌ الملف فارغ أو تالف."
        
        # تحويل إلى صور
        limit = max(1, int(limit_pages)) if limit_pages else 1
        dpi_val = max(150, min(300, int(dpi)))  # تحديد نطاق DPI
        
        logger.info(f"تحويل {limit} صفحة بدقة {dpi_val} DPI...")
        pages = pdf_to_images(pdf_bytes, dpi=dpi_val, max_pages=limit)
        
        if not pages:
            return "❌ لا توجد صفحات صالحة للمعالجة."
        
        # معالجة OCR
        results = []
        total_pages = len(pages)
        
        for i, (page_num, img) in enumerate(pages, 1):
            try:
                logger.info(f"معالجة الصفحة {page_num} ({i}/{total_pages})...")
                
                # OCR للصفحة
                text = ocr_page_gpu(img)
                
                if text and text.strip():
                    results.append(f"📄 **صفحة {page_num}**\n{'-'*50}\n{text}")
                else:
                    results.append(f"📄 **صفحة {page_num}**\n{'-'*50}\n⚠️ لم يتم استخراج نص من هذه الصفحة")
                
            except Exception as e:
                logger.error(f"خطأ في معالجة الصفحة {page_num}: {str(e)}")
                results.append(f"📄 **صفحة {page_num}**\n{'-'*50}\n❌ خطأ في المعالجة: {str(e)}")
        
        if not results:
            return "❌ لم يتم استخراج أي نص من الملف."
        
        final_result = "\n\n".join(results)
        logger.info(f"تمت معالجة {len(results)} صفحة بنجاح")
        
        return final_result
        
    except Exception as e:
        logger.error(f"خطأ عام في المعالجة: {str(e)}")
        logger.error(traceback.format_exc())
        return f"❌ حدث خطأ: {str(e)}"

# إنشاء واجهة Gradio
with gr.Blocks(
    title="Horoof OCR (H200 GPU)",
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        font-family: 'Arial', sans-serif;
    }
    .output-text {
        font-family: 'Courier New', monospace;
        line-height: 1.6;
    }
    """
) as demo:
    
    gr.Markdown("""
    # 🔤 Horoof OCR - استخراج النصوص العربية
    
    **النموذج**: Qwen2-VL-2B + Horoof LoRA  
    **المعالج**: H200 GPU  
    **الدعم**: PDF → نص عربي عالي الجودة
    
    ---
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            pdf_input = gr.File(
                label="📁 ارفع ملف PDF",
                file_types=[".pdf"],
                type="binary"
            )
            
            dpi_slider = gr.Slider(
                minimum=150,
                maximum=300,
                value=220,
                step=10,
                label="🎯 دقة التحويل (DPI)",
                info="دقة أعلى = جودة أفضل + وقت أطول"
            )
            
            pages_limit = gr.Number(
                value=2,
                minimum=1,
                maximum=10,
                precision=0,
                label="📊 عدد الصفحات المراد معالجتها",
                info="ابدأ بعدد قليل للاختبار"
            )
            
            process_btn = gr.Button(
                "🚀 بدء الاستخراج",
                variant="primary",
                size="lg"
            )
            
        with gr.Column(scale=2):
            output_text = gr.Textbox(
                label="📝 النص المستخرج",
                lines=25,
                max_lines=30,
                elem_classes=["output-text"],
                placeholder="سيظهر النص المستخرج هنا...",
                show_copy_button=True
            )
    
    gr.Markdown("""
    ---
    ### 💡 نصائح للحصول على أفضل النتائج:
    - **جودة الملف**: تأكد من أن PDF واضح وقابل للقراءة
    - **DPI**: استخدم 220-250 للنصوص العادية، 280-300 للخطوط الصغيرة
    - **عدد الصفحات**: ابدأ بصفحة أو اثنتين للاختبار
    - **أنواع النصوص**: يعمل بشكل ممتاز مع النصوص العربية المطبوعة والمكتوبة بوضوح
    """)
    
    # ربط الأحداث
    process_btn.click(
        fn=ocr_pdf,
        inputs=[pdf_input, dpi_slider, pages_limit],
        outputs=output_text,
        api_name="ocr_pdf"
    )

# إعداد الـ queue
demo.queue(
    concurrency_count=2,  # عدد العمليات المتزامنة
    max_size=10  # حد أقصى للطابور
)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        quiet=False
    )