import PyPDF2
import requests
import json
from reportlab.lib.pagesizes import letter
from reportlab.lib.units import inch
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.enums import TA_JUSTIFY
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
import os

# Configuration
PDF_PATH = "TaniaBorecMemoir(Ukr).pdf"  # Fichier original
OLLAMA_MODEL = "traductionUkrainienVersFrancais:latest"
OLLAMA_URL = "http://localhost:11434/api/generate"  # URL par défaut d'Ollama
OUTPUT_PDF_PATH = PDF_PATH.replace(".pdf", " (FR).pdf")  # Chemin du PDF de sortie

def extract_text_from_pdf(pdf_path):
    """Extrait le texte page par page d'un PDF sans les numéros de pages."""
    import re
    text_by_page = []
    with open(pdf_path, "rb") as file:
        reader = PyPDF2.PdfReader(file)
        for page in reader.pages:
            text = page.extract_text()
            # Supprime les numéros de pages (nombres seuls en début/fin de ligne)
            text = re.sub(r'^\d+\s*\n', '', text, flags=re.MULTILINE)
            text = re.sub(r'\n\s*\d+\s*$', '', text, flags=re.MULTILINE)
            text_by_page.append(text)
    return text_by_page

def merge_paragraphs_across_pages(pages_text):
    """Divise le texte en chunks raisonnables pour la traduction."""
    import re
    
    # Concatène tout le texte
    full_text = "\n".join(pages_text)
    
    # Essaie d'abord de diviser par les doubles sauts de ligne
    paragraphs = re.split(r'\n\s*\n+', full_text.strip())
    
    # Si on obtient seulement un paragraphe, on divise par une taille maximale
    if len(paragraphs) == 1:
        print("Aucune séparation par double saut de ligne détectée. Division par taille...")
        # Divise par les phrases (points suivis d'un espace)
        sentences = re.split(r'(?<=[.!?])\s+', full_text.strip())
        
        # Regroupe les phrases en chunks d'environ 1500 caractères
        max_chunk_size = 1500
        paragraphs = []
        current_chunk = ""
        
        for sentence in sentences:
            if len(current_chunk) + len(sentence) < max_chunk_size:
                current_chunk += (" " + sentence) if current_chunk else sentence
            else:
                if current_chunk:
                    paragraphs.append(current_chunk)
                current_chunk = sentence
        
        if current_chunk:
            paragraphs.append(current_chunk)
    else:
        # Normalise les sauts de ligne internes
        paragraphs = [re.sub(r'\n+', ' ', p.strip()) for p in paragraphs if p.strip()]
    
    return paragraphs

def send_to_ollama(prompt, model=OLLAMA_MODEL, context_size=128000):
    """Envoie une requête à Ollama et retourne la réponse."""
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "options": {"num_ctx": context_size}
    }
    response = requests.post(OLLAMA_URL, data=json.dumps(payload))
    if response.status_code == 200:
        return response.json()["response"]
    else:
        raise Exception(f"Erreur Ollama: {response.text}")

def register_unicode_font():
    """Enregistre une police TrueType qui supporte le cyrilique."""
    # Recherche une police système qui supporte le cyrilique
    font_paths = [
        r"C:\Windows\Fonts\DejaVuSans.ttf",
        r"C:\Windows\Fonts\Calibri.ttf",
        r"C:\Windows\Fonts\arial.ttf",
    ]
    
    for font_path in font_paths:
        if os.path.exists(font_path):
            try:
                pdfmetrics.registerFont(TTFont('UnicodeFont', font_path))
                return 'UnicodeFont'
            except Exception as e:
                print(f"Erreur lors de l'enregistrement de {font_path}: {e}")
    
    # Si aucune police spéciale trouvée, utilise Helvetica par défaut
    print("Aucune police Unicode trouvée, utilisation d'Helvetica")
    return 'Helvetica'

def create_pdf_from_results(results, output_path):
    """Crée un PDF à partir des résultats de traduction."""
    doc = SimpleDocTemplate(output_path, pagesize=letter, topMargin=inch, bottomMargin=inch)
    story = []
    
    # Enregistre une police qui supporte le cyrilique
    font_name = register_unicode_font()
    
    # Style personnalisé
    styles = getSampleStyleSheet()
    title_style = ParagraphStyle(
        'CustomTitle',
        parent=styles['Heading1'],
        fontSize=16,
        textColor='#1f4788',
        spaceAfter=0.3*inch,
        alignment=TA_JUSTIFY,
        fontName=font_name
    )
    
    page_style = ParagraphStyle(
        'PageHeading',
        parent=styles['Heading2'],
        fontSize=12,
        textColor='#1f4788',
        spaceAfter=0.2*inch,
        spaceBefore=0.2*inch,
        fontName=font_name
    )
    
    body_style = ParagraphStyle(
        'CustomBody',
        parent=styles['BodyText'],
        fontSize=11,
        alignment=TA_JUSTIFY,
        spaceAfter=0.2*inch,
        fontName=font_name
    )
    
    # Titre
    story.append(Paragraph("Traduction - Ukrainien vers Français", title_style))
    story.append(Spacer(1, 0.2*inch))
    
    # Contenu
    for page_num, translation in results.items():
        # Préserver la mise en page en convertissant les sauts de ligne
        formatted_text = translation.replace("\n", "<br/>")
        story.append(Paragraph(formatted_text, body_style))
        story.append(Spacer(1, 0.1*inch))
    
    # Construction du PDF
    doc.build(story)
    print(f"PDF généré avec succès : {output_path}")

def main():
    # Extraction du texte page par page
    pages = extract_text_from_pdf(PDF_PATH)
    print(f"Nombre de pages extraites : {len(pages)}")
    
    # Fusion des paragraphes qui s'étendent sur plusieurs pages
    paragraphs = merge_paragraphs_across_pages(pages)
    print(f"Nombre de paragraphes complets extraits : {len(paragraphs)}")

    # Dictionnaire pour stocker les résultats
    results = {}
    
    # Traitement des paragraphes complets
    for i, paragraph_text in enumerate(paragraphs, start=1):
        print(f"{15 * '-'} Traduction du paragraphe {i}/{len(paragraphs)}...\n{paragraph_text}\n")
        prompt = f"Traduis le texte suivant de l'ukrainien vers le français : {paragraph_text}"
        try:
            result = send_to_ollama(prompt)
            print(f"{result}.")
            results[i] = result
        except Exception as e:
            print(f"Erreur lors du traitement du paragraphe {i} : {e}")
            results[i] = f"Erreur lors du traitement du paragraphe {i} : {e}"
    
    # Création du PDF avec tous les résultats
    create_pdf_from_results(results, OUTPUT_PDF_PATH)

if __name__ == "__main__":
    main()