# -*- coding: utf-8 -*-
"""Untitled18.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1ebBGzEo4wbwwvReea_n0PRHdfYescKcs
"""

import os
import torch
from transformers import EncoderDecoderModel, AutoTokenizer
import re

# تعريف الثوابت
HF_REPO_ID = "bayan10/PuncAra-v1"

# متغيرات عامة
device = None
test_model = None
test_tokenizer = None

def initialize_model(repo_id=HF_REPO_ID):
    """
    تهيئة وإعداد كرت الشاشة وتحميل النموذج والـ Tokenizer من Hugging Face Hub.
    يتم استدعاء هذه الدالة مرة واحدة فقط في بداية تشغيل المشروع.
    """
    global device, test_model, test_tokenizer
    print(f"Loading test model directly from Hugging Face Hub: {repo_id}")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Loading test model to: {device}")

    if device.type == "cuda" and not torch.cuda.is_available():
        print("Warning: CUDA device requested, but torch.cuda.is_available() is False. Model will be loaded to CPU.")
        device = torch.device("cpu")

    test_model = EncoderDecoderModel.from_pretrained(repo_id).to(device)
    test_tokenizer = AutoTokenizer.from_pretrained(repo_id)

    # إعداد الـ Special tokens للـ Decoder والـ Encoder
    test_model.config.decoder_start_token_id = test_tokenizer.cls_token_id
    test_model.config.bos_token_id = test_tokenizer.cls_token_id
    test_model.config.eos_token_id = test_tokenizer.sep_token_id
    test_model.config.pad_token_id = test_tokenizer.pad_token_id
    print("Model and Tokenizer loaded successfully!")

def predict_chunk(text_chunk):
    """توليد التوقعات لعلامات الترقيم لقطعة نصية صغيرة لا تتعدى الـ 128 Token."""
    global device, test_model, test_tokenizer
    if test_model is None or test_tokenizer is None:
        raise RuntimeError("الموديل لم يتم تهيئته بعد. يرجى استدعاء initialize_model() أولاً.")

    # تطبيق الـ Preprocessing لتنظيف التشكيل قبل دخول النص للموديل
    text_chunk = arabic_preprocessing(text_chunk)

    inputs = test_tokenizer(text_chunk, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)

    outputs = test_model.generate(
          inputs.input_ids,
          attention_mask=inputs.attention_mask,
          decoder_start_token_id=test_tokenizer.cls_token_id,
          bos_token_id=test_tokenizer.cls_token_id,
          eos_token_id=test_tokenizer.sep_token_id,
          pad_token_id=test_tokenizer.pad_token_id,
          max_length=128,
          num_beams=3,
          repetition_penalty=1.2,
          length_penalty=1.0,
          early_stopping=True,
          do_sample=False
      )
    return test_tokenizer.decode(outputs[0], skip_special_tokens=True)

def arabic_preprocessing(text):
    """حذف الحركات التشكيلية لتوحيد المدخلات وتسهيل عمل الموديل."""
    arabic_diacritics = re.compile(r'[\u064B-\u0652]')
    return re.sub(arabic_diacritics, '', text).strip()

def arabic_postprocessing(text):
    """
    التنظيف والتحسين المطبعي وعلاج مشاكل دمج النصوص وعلامات الترقيم الزائدة.
    """
    if not text:
        return text

    # 1. حماية الأرقام والكسور والتوقيت من التحويل الخاطئ
    text = re.sub(r'(?<=\d),(?=\d)', '٪TEMP_COMMA٪', text)
    text = re.sub(r'(?<=\d):(?=\d)', '٪TEMP_COLON٪', text)

    # 2. التوحيد والتعريب المطبعي للعلامات
    text = text.replace(',', '،').replace(';', '؛').replace('?', '؟')

    # 3. ضبط المسافات الداخلية للأقواس وعلامات الاقتباس العربي
    text = re.sub(r'\(\s+', '(', text)
    text = re.sub(r'\s+\)', ')', text)
    text = re.sub(r'\[\s+', '[', text)
    text = re.sub(r'\s+\]', ']', text)
    text = re.sub(r'«\s+', '«', text)
    text = re.sub(r'\s+»', '»', text)

    # 4. منع تكرار العلامات الانفعالية عدا النقاط الثلاثية للحذف
    text = re.sub(r'([،؛:!؟])\1+', r'\1', text)
    text = re.sub(r'\.{4,}', '...', text)

    # 5. معالجة التناقضات المباشرة الناتجة عن تجميع الـ Chunks
    text = re.sub(r'[،؛:]+([.!؟])', r'\1', text)
    text = re.sub(r'،؛|؛،', '؛', text)
    text = re.sub(r'([!؟])\.', r'\1', text)

    # 6. مسح علامات الترقيم العشوائية إذا ظهرت أول النص
    text = re.sub(r'^[،؛:!؟. \t]+', '', text)

    # 7. ضمان مسافة فارغة واحدة بعد علامة الترقيم إذا تبعها كلام
    text = re.sub(r'([،؛:!؟.])(?=\S)', r'\1 ', text)

    # 8. إعادة الأرقام والكسور والتوقيت المحمية إلى أصلها
    text = text.replace('٪TEMP_COMMA٪', ',').replace('٪TEMP_COLON٪', ':')

    # 9. إلصاق علامات الترقيم بالكلمة السابقة لها مباشرة
    text = re.sub(r'\s+([،؛:!؟.])', r'\1', text)

    # 10. إزالة المسافات المتكررة الأفقية فقط (بدون لمس السطور الجديدة)
    text = re.sub(r'[ \t]+', ' ', text).strip()
    return text

def fix_punctuation(text):
    """معالجة الفقرة الواحدة الطويلة عبر تقسيمها لقطع غير متداخلة لمنع التكرار."""
    words = text.split()
    total_words = len(words)

    # جعل حجم الخطوة مساوياً لحجم النافذة يمنع تكرار الكلمات تماماً
    window_size = 50
    stride = 50

    if total_words <= window_size:
        result = predict_chunk(text)
    else:
        segments_output = []

        for i in range(0, total_words, stride):
            chunk_words = words[i : i + window_size]
            chunk_text = " ".join(chunk_words)

            if not chunk_text.strip():
                continue

            processed_segment = predict_chunk(chunk_text).strip()

            # مسح علامات الترقيم الناتجة عن القص الإجباري بين القطع
            is_last_segment = (i + window_size) >= total_words
            if not is_last_segment:
                punctuation_marks = ".?!،؛:؟!"
                if processed_segment and processed_segment[-1] in punctuation_marks:
                    # نمسح العلامة تماماً لأن السياق مستمر في القطعة اللي بعدها
                    processed_segment = processed_segment[:-1]

            segments_output.append(processed_segment)

        result = " ".join(segments_output)

    # تنظيف المسافات الزائدة والتكرار إن وجد
    result = re.sub(r'\s+', ' ', result).strip()
    return result

def process_full_document(text):
    if not text:
        return text

    # تقسيم بناءً على السطور الجديدة وتنظيف الأسطر الفارغة
    paragraphs = [p.strip() for p in text.split('\n') if p.strip()]
    processed_paragraphs = []

    for paragraph in paragraphs:
        # معالجة الفقرة المستقلة
        punctuated_paragraph = fix_punctuation(paragraph)
        cleaned_paragraph = arabic_postprocessing(punctuated_paragraph)
        processed_paragraphs.append(cleaned_paragraph)

    # الدمج بسطرين متباعدين لضمان الفصل البصري التام بين الفقرات
    return "\n\n".join(processed_paragraphs)