import os
import argostranslate.package
import argostranslate.translate
from tkinter import filedialog, Tk
import pdfplumber
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from PyPDF2 import PdfReader, PdfWriter
from io import BytesIO
from tqdm import tqdm
import re
import sys
import traceback
import fitz  # PyMuPDF pentru extragerea imaginilor și schemelor

# Limbile disponibile
SOURCE_LANGUAGES = {
    "1": ("en", "Engleză"),
    "2": ("fr", "Franceză"),
    "3": ("ru", "Rusă"),
    "4": ("el", "Greacă"),
    "5": ("grc", "Greacă veche"),
    "6": ("es", "Spaniolă"),
    "7": ("de", "Germană"),
    "8": ("it", "Italiană"),
    "9": ("la", "Latină")
}

TARGET_LANGUAGES = {
    "1": ("ro", "Română"),
    "2": ("en", "Engleză"),
    "3": ("fr", "Franceză")
}

# Configurare traducere Argos
def setup_argos_translation(from_lang, to_lang):
    """
    Configurează pachetul de traducere Argos pentru limbile specificate.
    Dacă pachetul direct nu există, încearcă pivotarea prin engleză.
    Returnează True dacă reușește, False dacă eșuează.
    """
    try:
        print(f"Se configurează traducerea {from_lang} -> {to_lang}...")
        print("Se actualizează indexul pachetelor de traducere...")
        with tqdm(total=1, desc="Actualizare index") as pbar:
            argostranslate.package.update_package_index()
            pbar.update(1)
        
        available_packages = argostranslate.package.get_available_packages()
        print(f"Perechi disponibile: {[(p.from_code, p.to_code) for p in available_packages]}")
        
        package = next(
            (p for p in available_packages if p.from_code == from_lang and p.to_code == to_lang),
            None
        )
        if package:
            installed_packages = argostranslate.package.get_installed_packages()
            if any(p.from_code == from_lang and p.to_code == to_lang for p in installed_packages):
                print(f"Pachetul {from_lang} -> {to_lang} este deja instalat.")
                return True
            
            print(f"Se descarcă pachetul de traducere {from_lang} -> {to_lang}...")
            with tqdm(total=1, desc="Descărcare pachet") as pbar:
                package_path = package.download()
                pbar.update(1)
            
            print(f"Se instalează pachetul {from_lang} -> {to_lang}...")
            with tqdm(total=1, desc="Instalare pachet") as pbar:
                argostranslate.package.install_from_path(package_path)
                pbar.update(1)
            
            print(f"Pachetul {from_lang} -> {to_lang} a fost configurat cu succes.")
            return True
        
        print(f"Nu există pachet direct pentru {from_lang} -> {to_lang}. Se încearcă pivotarea prin engleză...")
        if from_lang != "en" and to_lang != "en":
            package1 = next(
                (p for p in available_packages if p.from_code == from_lang and p.to_code == "en"),
                None
            )
            package2 = next(
                (p for p in available_packages if p.from_code == "en" and p.to_code == to_lang),
                None
            )
            if package1 and package2:
                print(f"Se configurează pivotarea: {from_lang} -> en -> {to_lang}")
                for pkg, pair in [(package1, (from_lang, "en")), (package2, ("en", to_lang))]:
                    installed_packages = argostranslate.package.get_installed_packages()
                    if any(p.from_code == pair[0] and p.to_code == pair[1] for p in installed_packages):
                        print(f"Pachetul {pair[0]} -> {pair[1]} este deja instalat.")
                        continue
                    print(f"Se descarcă pachetul de traducere {pair[0]} -> {pair[1]}...")
                    with tqdm(total=1, desc=f"Descărcare {pair[0]}->{pair[1]}") as pbar:
                        package_path = pkg.download()
                        pbar.update(1)
                    print(f"Se instalează pachetul {pair[0]} -> {pair[1]}...")
                    with tqdm(total=1, desc=f"Instalare {pair[0]}->{pair[1]}") as pbar:
                        argostranslate.package.install_from_path(package_path)
                        pbar.update(1)
                print(f"Pivotarea {from_lang} -> en -> {to_lang} a fost configurată cu succes.")
                return True
            else:
                print(f"Eroare: Nu se poate configura pivotarea. Lipsesc pachetele {from_lang} -> en sau en -> {to_lang}.")
                return False
        
        print(f"Eroare: Nu există pachet de traducere disponibil pentru {from_lang} -> {to_lang} și pivotarea nu este posibilă.")
        return False
    except Exception as e:
        print(f"Eroare la configurarea traducătorului {from_lang} -> {to_lang}: {str(e)}")
        print("Detalii eroare:")
        traceback.print_exc()
        return False

# Selectare limbă
def select_language(prompt, languages):
    """
    Afișează limbile disponibile și cere utilizatorului să aleagă un număr.
    Returnează codul limbii selectate sau None dacă selecția este invalidă.
    """
    print(prompt)
    for num, (_, name) in languages.items():
        print(f"{num}: {name}")
    choice = input("Introdu numărul limbii: ").strip()
    if choice in languages:
        return languages[choice][0]
    print("Selecție invalidă.")
    return None

# Extragere text și proprietăți din PDF
def extract_text_with_properties(pdf_file):
    """
    Extrage textul din PDF linie cu linie, cu pozițiile și proprietățile sale (font, dimensiune, stil).
    Returnează o listă de dicționare cu text, coordonate și proprietăți sau None dacă eșuează.
    """
    text_lines = []
    try:
        with pdfplumber.open(pdf_file) as pdf:
            total_pages = len(pdf.pages)
            print(f"Se procesează {total_pages} pagini...")
            with tqdm(total=total_pages, desc="Extragere text") as pbar:
                for page_num, page in enumerate(pdf.pages, 1):
                    chars = page.chars  # Extragem caracterele cu proprietăți
                    current_line = []
                    last_y = None
                    for char in chars:
                        y = char["top"]
                        if last_y is None or abs(y - last_y) > 2:  # Grupăm caracterele în linii
                            if current_line:
                                text = "".join(c["text"] for c in current_line)
                                fontname = current_line[0]["fontname"]
                                size = current_line[0]["size"]
                                x0 = min(c["x0"] for c in current_line)
                                x1 = max(c["x1"] for c in current_line)
                                top = current_line[0]["top"]
                                text_lines.append({
                                    "text": text,
                                    "x0": x0,
                                    "top": top,
                                    "x1": x1,
                                    "page": page_num,
                                    "width": page.width,
                                    "height": page.height,
                                    "fontname": fontname,
                                    "size": size
                                })
                            current_line = [char]
                        else:
                            current_line.append(char)
                        last_y = y
                    if current_line:  # Procesăm ultima linie
                        text = "".join(c["text"] for c in current_line)
                        fontname = current_line[0]["fontname"]
                        size = current_line[0]["size"]
                        x0 = min(c["x0"] for c in current_line)
                        x1 = max(c["x1"] for c in current_line)
                        top = current_line[0]["top"]
                        text_lines.append({
                            "text": text,
                            "x0": x0,
                            "top": top,
                            "x1": x1,
                            "page": page_num,
                            "width": page.width,
                            "height": page.height,
                            "fontname": fontname,
                            "size": size
                        })
                    pbar.update(1)
                    print(f"Pagina {page_num}/{total_pages} procesată.")
        return text_lines
    except Exception as e:
        print(f"Eroare la extragerea textului din PDF: {str(e)}")
        return None

# Detectare și gestionare formule matematice
def handle_math_formulas(text):
    """
    Detectează formulele matematice între $...$ sau $$...$$ și le separă.
    Returnează o pereche: lista formulelor și textul cu placeholder.
    """
    math_pattern = r'\$\$?.*?\$\$?|\$.*?\$'
    formulas = re.findall(math_pattern, text)
    text_with_placeholder = re.sub(math_pattern, "###MATH###", text)
    return formulas, text_with_placeholder

# Traducere text
def translate_with_formulas(text, from_lang, to_lang):
    """
    Traduce textul păstrând formulele matematice intacte.
    Returnează textul tradus cu formulele reintroduse.
    """
    try:
        formulas, text_placeholder = handle_math_formulas(text)
        translated = argostranslate.translate.translate(text_placeholder, from_lang, to_lang)
        for formula in formulas:
            translated = translated.replace("###MATH###", formula, 1)
        return translated
    except Exception as e:
        print(f"Eroare la traducerea textului: {str(e)}")
        return text

# Generare PDF tradus cu imagini și formatare
def generate_translated_pdf(text_data, translated_texts, input_pdf, output_file):
    """
    Creează un PDF nou cu textul tradus, păstrând imaginile, schemele și formatarea originală.
    """
    try:
        # Încărcăm fontul implicit
        pdfmetrics.registerFont(TTFont("DejaVuSans", "DejaVuSans.ttf"))
        pdfmetrics.registerFont(TTFont("DejaVuSans-Bold", "DejaVuSans-Bold.ttf"))
        pdfmetrics.registerFont(TTFont("DejaVuSans-Oblique", "DejaVuSans-Oblique.ttf"))

        # Citim PDF-ul original
        pdf_reader = PdfReader(input_pdf)
        pdf_writer = PdfWriter()

        print("Se generează PDF-ul tradus...")
        with tqdm(total=len(text_data), desc="Generare PDF") as pbar:
            current_page = 1
            temp_pdf = BytesIO()
            pdf_canvas = canvas.Canvas(temp_pdf, pagesize=letter)

            for i, item in enumerate(text_data):
                if item["page"] != current_page:
                    pdf_canvas.showPage()
                    current_page = item["page"]
                    pdf_canvas.setPageSize((item["width"], item["height"]))

                translated_text = translated_texts[i]
                y_pos = item["height"] - item["top"]
                
                # Determinăm fontul și stilul
                fontname = "DejaVuSans"
                if "Bold" in item["fontname"]:
                    fontname = "DejaVuSans-Bold"
                elif "Italic" in item["fontname"] or "Oblique" in item["fontname"]:
                    fontname = "DejaVuSans-Oblique"
                
                font_size = item["size"]
                pdf_canvas.setFont(fontname, font_size)

                # Ajustăm dimensiunea fontului dacă textul e prea lung
                original_width = item["x1"] - item["x0"]
                text_width = pdf_canvas.stringWidth(translated_text, fontname, font_size)
                if text_width > original_width:
                    font_size = font_size * (original_width / text_width)
                    font_size = max(font_size, 6)  # Minim 6pt
                    pdf_canvas.setFont(fontname, font_size)

                pdf_canvas.drawString(item["x0"], y_pos, translated_text)
                pbar.update(1)

            pdf_canvas.save()
            temp_pdf.seek(0)

            # Combinăm PDF-ul original cu textul tradus
            translated_pdf = PdfReader(temp_pdf)
            for page_num in range(len(pdf_reader.pages)):
                original_page = pdf_reader.pages[page_num]
                translated_page = translated_pdf.pages[page_num] if page_num < len(translated_pdf.pages) else None
                new_page = original_page
                if translated_page:
                    new_page.merge_page(translated_page)  # Suprapunem textul tradus
                pdf_writer.add_page(new_page)

            # Salvăm PDF-ul final
            with open(output_file, "wb") as f:
                pdf_writer.write(f)
            print(f"PDF generat cu succes: {output_file}")
    except Exception as e:
        print(f"Eroare la generarea PDF-ului: {str(e)}")
        raise

# Funcția principală
def process_pdf_translation():
    """Procesează și traduce un PDF selectat de utilizator."""
    try:
        # Selectăm limba sursă
        from_lang = select_language("Selectați limba sursă:", SOURCE_LANGUAGES)
        if not from_lang:
            print("Eroare: Nicio limbă sursă selectată.")
            return

        # Selectăm limba țintă
        to_lang = select_language("Selectați limba țintă:", TARGET_LANGUAGES)
        if not to_lang:
            print("Eroare: Nicio limbă țintă selectată.")
            return

        # Verificăm traducerea Argos
        if not setup_argos_translation(from_lang, to_lang):
            print(f"Eroare: Nu s-a putut configura traducerea {from_lang} -> {to_lang}.")
            return

        # Selectăm fișierul PDF
        root = Tk()
        root.withdraw()
        pdf_file = filedialog.askopenfilename(
            title="Selectează fișierul PDF de tradus",
            filetypes=[("PDF files", "*.pdf")]
        )
        root.destroy()

        if not pdf_file:
            print("Niciun fișier PDF selectat.")
            return

        # Setăm folderul de ieșire
        script_dir = os.path.dirname(os.path.abspath(__file__))
        output_dir = os.path.join(script_dir, "translated")
        os.makedirs(output_dir, exist_ok=True)

        # Extragem textul și proprietățile
        text_data = extract_text_with_properties(pdf_file)
        if not text_data:
            print("Eroare: Nu s-a putut extrage textul din PDF.")
            return

        # Traducem textul
        original_texts = [item["text"] for item in text_data]
        translated_texts = []
        print("Se traduce textul...")
        with tqdm(total=len(original_texts), desc="Traducere") as pbar:
            for text in original_texts:
                translated_texts.append(translate_with_formulas(text, from_lang, to_lang))
                pbar.update(1)

        # Generăm PDF-ul tradus
        output_file = os.path.join(
            output_dir, f"{os.path.splitext(os.path.basename(pdf_file))[0]}_translated_{to_lang}.pdf"
        )
        generate_translated_pdf(text_data, translated_texts, pdf_file, output_file)

    except Exception as e:
        print(f"Eroare generală în procesare: {str(e)}")
        print("Detalii eroare:")
        traceback.print_exc()

if __name__ == "__main__":
    process_pdf_translation()