from transformers import pipeline
import fitz  # PyMuPDF
from tqdm import tqdm  # For progress bars
import torch  # For device management
import re  # For text formatting

class TranslatorAR:
    """
    TranslatorAR is a class for translating text between Arabic, English, and French using Hugging Face's translation models.
    It supports both direct text translation and PDF file translation.
    """
    def __init__(self):
        """
        Initializes the TranslatorAR with translation pipelines and GPU/CPU configuration.
        """
        self.device = 0 if torch.cuda.is_available() else -1
        self.translators = self._initialize_pipelines()
        self.max_length = 512

    def _initialize_pipelines(self):
        """
        Initializes translation pipelines for supported language pairs.

        :return: Dictionary mapping language pairs to translation pipelines
        """
        models = {
            ("ar", "en"): "Helsinki-NLP/opus-mt-ar-en",
            ("en", "ar"): "Helsinki-NLP/opus-mt-en-ar",
            ("fr", "en"): "Helsinki-NLP/opus-mt-fr-en",
            ("en", "fr"): "Helsinki-NLP/opus-mt-en-fr",
            ("ar", "fr"): "Helsinki-NLP/opus-mt-ar-fr",
            ("fr", "ar"): "Helsinki-NLP/opus-mt-fr-ar"
        }
        return {pair: pipeline("translation", model=model, device=self.device) for pair, model in models.items()}

    def translate(self, text, src_lang, tgt_lang):
        """
        Translates text from source language to target language.

        :param text: Text to translate
        :param src_lang: Source language code
        :param tgt_lang: Target language code
        :return: Translated text with formatting
        :raises ValueError: If the language pair is not supported
        """
        self._validate_language_pair(src_lang, tgt_lang)
        chunks = self._chunk_text(text)
        translated_chunks = [self._translate_chunk(chunk, src_lang, tgt_lang) for chunk in tqdm(chunks, desc="Translating chunks")]
        return self._format_text(' '.join(translated_chunks))

    def _validate_language_pair(self, src_lang, tgt_lang):
        """
        Validates if the given language pair is supported.

        :param src_lang: Source language code
        :param tgt_lang: Target language code
        :raises ValueError: If the language pair is not supported
        """
        if (src_lang, tgt_lang) not in self.translators:
            raise ValueError(f"Translation from {src_lang} to {tgt_lang} is not supported.")

    def _chunk_text(self, text):
        """
        Splits text into manageable chunks based on maximum length.

        :param text: Text to chunk
        :return: List of text chunks
        """
        return [text[i:i+self.max_length] for i in range(0, len(text), self.max_length)]

    def _translate_chunk(self, chunk, src_lang, tgt_lang):
        """
        Translates a single chunk of text.

        :param chunk: Text chunk to translate
        :param src_lang: Source language code
        :param tgt_lang: Target language code
        :return: Translated text chunk
        """
        translator = self.translators[(src_lang, tgt_lang)]
        return translator(chunk)[0]['translation_text']

    def _format_text(self, text):
        """
        Formats the text with sentence and line breaks.

        :param text: Text to format
        :return: Formatted text
        """
        text = re.sub(r'(?<!\n)\.\s', '.\n', text)
        text = re.sub(r'(?<!\n)\!\s', '!\n', text)
        text = re.sub(r'(?<!\n)\?\s', '?\n', text)
        return re.sub(r'\n+', '\n', text).strip()

    def translate_pdf(self, pdf_path, src_lang, tgt_lang, output_file):
        """
        Extracts text from a PDF, translates it, and saves the translated text to a file.

        :param pdf_path: Path to the PDF file
        :param src_lang: Source language code
        :param tgt_lang: Target language code
        :param output_file: Path to the output file
        :return: Tuple of extracted text and translated text
        """
        extracted_text = self._extract_text_from_pdf(pdf_path)
        translated_text = self.translate(extracted_text, src_lang, tgt_lang)
        self._save_to_file(translated_text, output_file)
        return extracted_text, translated_text

    def _extract_text_from_pdf(self, pdf_path):
        """
        Extracts text from each page of a PDF file.

        :param pdf_path: Path to the PDF file
        :return: Extracted text
        """
        document = fitz.open(pdf_path)
        text = ''.join(page.get_text() for page in tqdm(document, desc="Extracting pages"))
        return text

    def _save_to_file(self, text, file_path):
        """
        Saves the provided text to a file.

        :param text: Text to save
        :param file_path: Path to the file
        """
        with open(file_path, 'w', encoding='utf-8') as file:
            file.write(self._format_text(text))

    def print_translation(self, src_text, translated_text, src_lang, tgt_lang):
        """
        Prints the source and translated text with formatting.

        :param src_text: Source text
        :param translated_text: Translated text
        :param src_lang: Source language code
        :param tgt_lang: Target language code
        """
        print(f"--- Translation from {src_lang} to {tgt_lang} ---")
        print(f"Source Text ({src_lang}):\n{self._format_text(src_text)}\n")
        print(f"Translated Text ({tgt_lang}):\n{self._format_text(translated_text)}\n")

    def print_pdf_translation_results(self, pdf_path, src_lang, tgt_lang, output_file):
        """
        Prints the results of translating a PDF file.

        :param pdf_path: Path to the PDF file
        :param src_lang: Source language code
        :param tgt_lang: Target language code
        :param output_file: Path to the output file
        """
        extracted_text, translated_text = self.translate_pdf(pdf_path, src_lang, tgt_lang, output_file)
        print(f"--- PDF Translation Results ---")
        print(f"Extracted Text from PDF:\n{self._format_text(extracted_text)}\n")
        print(f"Translated Text saved to '{output_file}':\n{self._format_text(translated_text)}\n")

# Example usage
if __name__ == "__main__":
    translator = TranslatorAR()

    # Path to the PDF file
    PDF_FILE_PATH = "arabic-text.pdf"

    # Example 1: Arabic to English
    text_ar_en = "هذا المساق مقدم من Hugging Face."
    translated_text_ar_en = translator.translate(text_ar_en, "ar", "en")
    translator.print_translation(text_ar_en, translated_text_ar_en, "ar", "en")

    # Example 2: English to Arabic
    text_en_ar = "This course is produced by Hugging Face."
    translated_text_en_ar = translator.translate(text_en_ar, "en", "ar")
    translator.print_translation(text_en_ar, translated_text_en_ar, "en", "ar")

    # Example 3: Arabic to French
    text_ar_fr = "أهلاً وسهلاً بكم في هذا المساق."
    translated_text_ar_fr = translator.translate(text_ar_fr, "ar", "fr")
    translator.print_translation(text_ar_fr, translated_text_ar_fr, "ar", "fr")

    # Example 4: French to Arabic
    text_fr_ar = "Bienvenue dans ce cours."
    translated_text_fr_ar = translator.translate(text_fr_ar, "fr", "ar")
    translator.print_translation(text_fr_ar, translated_text_fr_ar, "fr", "ar")

    # Example 1.bis: Translate Arabic PDF to English
    translator.print_pdf_translation_results(PDF_FILE_PATH, "ar", "en", 'translated_to_english.txt')

    # Example 2.bis: Translate Arabic PDF to French
    translator.print_pdf_translation_results(PDF_FILE_PATH, "ar", "fr", 'translated_to_french.txt')
