import gradio as gr
import easyocr
from transformers import pipeline
import numpy as np
from PIL import Image
import torch
import re
from datetime import datetime
import PyPDF2
import io
import fitz  # PyMuPDF for better PDF handling
import os

# Set cache directory for transformers
os.environ["TRANSFORMERS_CACHE"] = "/tmp/transformers_cache"

# Check GPU availability
device = 0 if torch.cuda.is_available() else -1
gpu_available = torch.cuda.is_available()

# Initialize EasyOCR reader with multiple languages
reader = easyocr.Reader(['en', 'fr', 'es', 'de'], gpu=gpu_available)

# Initialize summarization pipeline
summarizer = pipeline(
    "summarization",
    model="facebook/bart-large-cnn",
    device=device
)

def extract_metadata(text):
    """
    Extract article metadata from text using pattern matching
    """
    metadata = {
        "title": "Non détecté",
        "author": "Non détecté",
        "date": "Non détecté",
        "location": "Non détecté"
    }
    
    lines = [line.strip() for line in text.split('\n') if line.strip()]
    
    # Extract Title (usually the first long line)
    for line in lines[:10]:
        if len(line) > 20 and len(line) < 200:
            metadata["title"] = line
            break
    
    # Extract Author - Common patterns
    author_patterns = [
        r'(?:by|par|por|von|de)\s+([A-Z][a-z]+(?:\s+[A-Z][a-z]+){1,3})',
        r'(?:author|auteur|autor|écrit par):\s*([A-Z][a-z]+(?:\s+[A-Z][a-z]+){1,3})',
        r'^([A-Z][a-z]+\s+[A-Z][a-z]+(?:\s+[A-Z][a-z]+)?)(?:\s*[-–—]|\s*\|)',
    ]
    
    for pattern in author_patterns:
        match = re.search(pattern, text, re.MULTILINE | re.IGNORECASE)
        if match:
            metadata["author"] = match.group(1).strip()
            break
    
    # Extract Date
    date_patterns = [
        r'\b(\d{1,2}[\s/\-\.]\w+[\s/\-\.]\d{2,4})\b',
        r'\b(\w+\s+\d{1,2},?\s+\d{4})\b',
        r'\b(\d{4}[\s/\-\.]\d{1,2}[\s/\-\.]\d{1,2})\b',
        r'\b(\d{1,2}[\s/\-\.]\d{1,2}[\s/\-\.]\d{2,4})\b',
    ]
    
    for pattern in date_patterns:
        match = re.search(pattern, text)
        if match:
            metadata["date"] = match.group(1).strip()
            break
    
    # Extract Location
    location_patterns = [
        r'\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)?),\s*([A-Z]{2,})\b',
        r'(?:in|à|en|in)\s+([A-Z][a-z]+(?:,?\s+[A-Z][a-z]+)?)',
    ]
    
    for pattern in location_patterns:
        match = re.search(pattern, text)
        if match:
            if len(match.groups()) > 1:
                metadata["location"] = f"{match.group(1)}, {match.group(2)}"
            else:
                metadata["location"] = match.group(1).strip()
            break
    
    return metadata

def extract_text_from_pdf(pdf_file):
    """
    Extract text from PDF file using PyMuPDF
    """
    try:
        pdf_bytes = pdf_file if isinstance(pdf_file, bytes) else pdf_file.read()
        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
        
        full_text = ""
        page_texts = []
        
        for page_num in range(len(doc)):
            page = doc[page_num]
            text = page.get_text()
            page_texts.append({
                "page_number": page_num + 1,
                "text": text,
                "word_count": len(text.split())
            })
            full_text += f"\n\n--- Page {page_num + 1} ---\n\n{text}"
        
        doc.close()
        return full_text, page_texts, len(doc)
    except Exception as e:
        raise Exception(f"Erreur lors de la lecture du PDF: {str(e)}")

def extract_text_from_images(image_files, progress=gr.Progress()):
    """
    Extract text from multiple images using OCR
    """
    full_text = ""
    page_texts = []
    
    for idx, image_file in enumerate(progress.tqdm(image_files, desc="Traitement des images")):
        try:
            if isinstance(image_file, str):
                image = Image.open(image_file)
            else:
                image = Image.open(image_file.name) if hasattr(image_file, 'name') else Image.open(image_file)
            
            image_array = np.array(image)
            results = reader.readtext(image_array, detail=1)
            
            page_text = " ".join([result[1] for result in results])
            page_texts.append({
                "page_number": idx + 1,
                "text": page_text,
                "word_count": len(page_text.split())
            })
            
            full_text += f"\n\n--- Page {idx + 1} ---\n\n{page_text}"
        except Exception as e:
            page_texts.append({
                "page_number": idx + 1,
                "text": f"Erreur: {str(e)}",
                "word_count": 0
            })
    
    return full_text, page_texts, len(image_files)

def chunk_text_for_summary(text, chunk_size=800):
    """
    Split text into chunks for better summarization of long documents
    """
    words = text.split()
    chunks = []
    for i in range(0, len(words), chunk_size):
        chunk = " ".join(words[i:i + chunk_size])
        chunks.append(chunk)
    return chunks

def generate_hierarchical_summary(text, min_length, max_length):
    """
    Generate a comprehensive summary using hierarchical approach for long documents
    """
    words = text.split()
    word_count = len(words)
    
    if word_count < 30:
        return "⚠️ Texte trop court pour générer un résumé (minimum 30 mots requis)."
    
    try:
        # For very long documents (>3000 words), use multi-stage summarization
        if word_count > 3000:
            # Stage 1: Split into chunks and summarize each
            chunks = chunk_text_for_summary(text, chunk_size=1000)
            chunk_summaries = []
            
            for chunk in chunks[:5]:  # Limit to first 5 chunks for performance
                try:
                    summary = summarizer(
                        chunk,
                        max_length=150,
                        min_length=50,
                        do_sample=False,
                        truncation=True
                    )
                    chunk_summaries.append(summary[0]['summary_text'])
                except:
                    continue
            
            # Stage 2: Combine and summarize the summaries
            combined_summary = " ".join(chunk_summaries)
            if len(combined_summary.split()) > 100:
                final_summary = summarizer(
                    combined_summary,
                    max_length=max_length,
                    min_length=min_length,
                    do_sample=False,
                    truncation=True
                )
                return final_summary[0]['summary_text']
            else:
                return combined_summary
        
        # For medium documents (1000-3000 words)
        elif word_count > 1000:
            text_to_summarize = " ".join(words[:1500])
            max_len = min(max_length, 300)
            min_len = min(min_length, 80)
        
        # For short documents (<1000 words)
        else:
            text_to_summarize = text
            max_len = min(max_length, word_count)
            min_len = min(min_length, word_count // 3)
        
        summary = summarizer(
            text_to_summarize,
            max_length=max_len,
            min_length=min_len,
            do_sample=False,
            truncation=True
        )
        return summary[0]['summary_text']
    
    except Exception as e:
        return f"Erreur lors de la génération du résumé: {str(e)}"

def process_document(file_input, file_type, min_summary_length, max_summary_length, progress=gr.Progress()):
    """
    Process uploaded document (PDF or images) and extract all information
    """
    if not file_input:
        return "Veuillez télécharger un fichier.", "", "", "", "", "", "", "N/A", 0, 0
    
    try:
        progress(0, desc="🚀 Démarrage de l'extraction...")
        
        # Extract text based on file type
        if file_type == "PDF":
            progress(0.1, desc="📄 Lecture du PDF...")
            full_text, page_texts, num_pages = extract_text_from_pdf(file_input)
            progress(0.4, desc=f"✅ PDF extrait: {num_pages} pages")
        else:  # Images
            progress(0.1, desc="🖼️ Préparation des images...")
            files = file_input if isinstance(file_input, list) else [file_input]
            full_text, page_texts, num_pages = extract_text_from_images(files, progress)
            progress(0.4, desc=f"✅ {num_pages} images traitées")
        
        if not full_text.strip():
            return "Aucun texte détecté dans le document.", "", "", "", "", "", "", "N/A", 0, num_pages
        
        progress(0.5, desc="🔍 Extraction des métadonnées...")
        
        # Extract metadata
        metadata = extract_metadata(full_text)
        
        # Word count
        total_words = sum([p["word_count"] for p in page_texts])
        
        progress(0.6, desc="📊 Analyse des pages...")
        
        # Create detailed page summary
        page_summary = "📈 STATISTIQUES PAR PAGE:\n\n"
        page_summary += "\n".join([
            f"📄 Page {p['page_number']}: {p['word_count']} mots"
            for p in page_texts[:15]  # Show first 15 pages
        ])
        
        if len(page_texts) > 15:
            remaining_pages = len(page_texts) - 15
            remaining_words = sum([p['word_count'] for p in page_texts[15:]])
            page_summary += f"\n\n📚 ... et {remaining_pages} pages supplémentaires ({remaining_words} mots)"
        
        page_summary += f"\n\n📊 TOTAL: {num_pages} pages | {total_words} mots"
        
        progress(0.7, desc="🤖 Génération du résumé intelligent...")
        
        # Generate comprehensive summary
        summary_text = generate_hierarchical_summary(full_text, min_summary_length, max_summary_length)
        
        progress(0.9, desc="✨ Finalisation...")
        
        # Calculate confidence (for OCR only)
        confidence = "N/A (PDF)" if file_type == "PDF" else "~85%"
        
        progress(1.0, desc="✅ Terminé!")
        
        return (
            full_text,
            metadata["title"],
            metadata["author"],
            metadata["date"],
            metadata["location"],
            summary_text,
            page_summary,
            confidence,
            total_words,
            num_pages
        )
    
    except Exception as e:
        return f"❌ Erreur: {str(e)}", "", "", "", "", "", "", "N/A", 0, 0

def clear_all():
    """Clear all inputs and outputs"""
    return None, "", "", "", "", "", "", "", "N/A", 0, 0

def update_summary_lengths(style):
    """Update summary length sliders based on style"""
    if style == "Concis":
        return 30, 150
    elif style == "Équilibré":
        return 50, 250
    else:  # Détaillé
        return 80, 400

def export_results(full_text, title, author, date, location, summary, page_info):
    """Export results to a formatted text file"""
    if not full_text or full_text.startswith("Veuillez") or full_text.startswith("Aucun") or full_text.startswith("❌"):
        return None
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"/tmp/document_analysis_{timestamp}.txt"
    
    content = f"""
╔══════════════════════════════════════════════════════════════╗
║        ANALYSE DE DOCUMENT - EXTRACTION DE DONNÉES           ║
╚══════════════════════════════════════════════════════════════╝

📰 TITRE:
{title}

✍️ AUTEUR:
{author}

📅 DATE:
{date}

📍 LOCALISATION:
{location}

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

{page_info}

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

📝 RÉSUMÉ INTELLIGENT:

{summary}

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

📄 TEXTE COMPLET EXTRAIT:

{full_text}

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Généré le: {datetime.now().strftime("%d/%m/%Y à %H:%M:%S")}
Par: Extracteur de Documents Multi-Pages v2.0
"""
    
    try:
        # Save to temporary file
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(content)
        return filename
    except Exception as e:
        print(f"Erreur export: {e}")
        return None

# Custom CSS
custom_css = """
.gradio-container {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}
.upload-area {
    border: 2px dashed #667eea;
    border-radius: 10px;
    padding: 20px;
    background: #f8f9ff;
}
.stat-box {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    padding: 12px;
    border-radius: 8px;
    color: white;
    text-align: center;
    font-weight: bold;
}
footer {visibility: hidden}
"""

# Create Gradio interface
with gr.Blocks(title="Extracteur de Documents Multi-Pages", css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 📚 Extracteur de Documents Multi-Pages
        ### Analyse complète de livres, articles et documents (PDF ou images)
        
        **Formats supportés:**
        - 📄 PDF (extraction de texte native)
        - 🖼️ Images multiples (JPG, PNG) avec OCR
        
        **Informations extraites:**
        - 📰 Titre du document
        - ✍️ Auteur
        - 📅 Date de publication
        - 📍 Localisation
        - 📝 Résumé automatique complet
        - 📊 Statistiques par page
        """
    )
    
    with gr.Row():
        gr.Markdown(f"**Statut:** {'🟢 GPU Activé' if gpu_available else '🔵 Mode CPU'}")
    
    with gr.Tabs() as tabs:
        with gr.TabItem("📄 Upload PDF"):
            with gr.Row():
                with gr.Column(scale=1):
                    pdf_input = gr.File(
                        label="📤 Télécharger un fichier PDF",
                        file_types=[".pdf"],
                        type="binary"
                    )
                    
                    with gr.Accordion("⚙️ Paramètres de résumé", open=True):
                        gr.Markdown("**Contrôlez la longueur de votre résumé:**")
                        min_summary_pdf = gr.Slider(
                            minimum=20,
                            maximum=150,
                            value=50,
                            step=10,
                            label="📏 Longueur minimale (mots)",
                            info="Plus court = résumé concis"
                        )
                        max_summary_pdf = gr.Slider(
                            minimum=100,
                            maximum=500,
                            value=250,
                            step=25,
                            label="📏 Longueur maximale (mots)",
                            info="Plus long = résumé détaillé"
                        )
                        
                        summary_style = gr.Radio(
                            choices=["Concis", "Équilibré", "Détaillé"],
                            value="Équilibré",
                            label="📝 Style de résumé"
                        )
                        
                        summary_style.change(
                            fn=update_summary_lengths,
                            inputs=[summary_style],
                            outputs=[min_summary_pdf, max_summary_pdf]
                        )
                    
                    analyze_pdf_btn = gr.Button("🔍 Analyser le PDF", variant="primary", size="lg")
        
        with gr.TabItem("🖼️ Upload Images"):
            with gr.Row():
                with gr.Column(scale=1):
                    images_input = gr.File(
                        label="📤 Télécharger plusieurs images (pages du livre)",
                        file_types=["image"],
                        file_count="multiple"
                    )
                    
                    gr.Markdown(
                        """
                        💡 **Conseil:** Nommez vos fichiers dans l'ordre (page1.jpg, page2.jpg, etc.)
                        pour une extraction séquentielle correcte.
                        """
                    )
                    
                    with gr.Accordion("⚙️ Paramètres de résumé", open=True):
                        gr.Markdown("**Contrôlez la longueur de votre résumé:**")
                        min_summary_img = gr.Slider(
                            minimum=20,
                            maximum=150,
                            value=50,
                            step=10,
                            label="📏 Longueur minimale (mots)",
                            info="Plus court = résumé concis"
                        )
                        max_summary_img = gr.Slider(
                            minimum=100,
                            maximum=500,
                            value=250,
                            step=25,
                            label="📏 Longueur maximale (mots)",
                            info="Plus long = résumé détaillé"
                        )
                        
                        summary_style_img = gr.Radio(
                            choices=["Concis", "Équilibré", "Détaillé"],
                            value="Équilibré",
                            label="📝 Style de résumé"
                        )
                        
                        summary_style_img.change(
                            fn=update_summary_lengths,
                            inputs=[summary_style_img],
                            outputs=[min_summary_img, max_summary_img]
                        )
                    
                    analyze_images_btn = gr.Button("🔍 Analyser les Images", variant="primary", size="lg")
    
    # Statistics Row
    with gr.Row():
        confidence_output = gr.Textbox(
            label="📊 Confiance",
            value="N/A",
            interactive=False,
            scale=1
        )
        word_count_output = gr.Textbox(
            label="📝 Total Mots",
            value="0",
            interactive=False,
            scale=1
        )
        page_count_output = gr.Textbox(
            label="📄 Nombre Pages",
            value="0",
            interactive=False,
            scale=1
        )
    
    # Metadata Section
    gr.Markdown("### 📋 Métadonnées Extraites")
    
    with gr.Row():
        title_output = gr.Textbox(
            label="📰 Titre",
            placeholder="Le titre sera extrait automatiquement...",
            lines=2,
            show_copy_button=True,
            scale=2
        )
        author_output = gr.Textbox(
            label="✍️ Auteur",
            placeholder="L'auteur sera détecté...",
            show_copy_button=True,
            scale=1
        )
    
    with gr.Row():
        date_output = gr.Textbox(
            label="📅 Date",
            placeholder="Date...",
            show_copy_button=True,
            scale=1
        )
        location_output = gr.Textbox(
            label="📍 Localisation",
            placeholder="Lieu...",
            show_copy_button=True,
            scale=1
        )
    
    # Results Section
    with gr.Row():
        with gr.Column():
            summary_output = gr.Textbox(
                label="📝 Résumé du Document",
                lines=8,
                placeholder="Le résumé sera généré automatiquement...",
                show_copy_button=True
            )
            
            page_info_output = gr.Textbox(
                label="📊 Informations par Page",
                lines=6,
                placeholder="Statistiques par page...",
                show_copy_button=True
            )
        
        with gr.Column():
            extracted_output = gr.Textbox(
                label="📄 Texte Complet Extrait",
                lines=14,
                placeholder="Le texte extrait apparaîtra ici...",
                show_copy_button=True
            )
    
    # Export Section
    with gr.Row():
        clear_btn = gr.Button("🗑️ Effacer Tout", variant="secondary")
        export_btn = gr.Button("💾 Exporter les Résultats", variant="secondary")
    
    export_output = gr.File(label="📥 Fichier d'Export")
    
    gr.Markdown(
        """
        ---
        ### 💡 Conseils pour de meilleurs résultats:
        
        **Pour les PDF:**
        - ✅ Utilisez des PDF avec texte sélectionnable (pas des scans)
        - ✅ Les PDF natifs donnent de meilleurs résultats que les PDF scannés
        - ⚡ Traitement ultra-rapide (quelques secondes pour 100+ pages)
        
        **Pour les Images:**
        - ✅ Images haute résolution (min. 1200x1600 pixels)
        - ✅ Bon éclairage et contraste
        - ✅ Texte bien visible et lisible
        - ✅ Nommez les fichiers dans l'ordre (page1.jpg, page2.jpg...)
        - ⏱️ Comptez ~2-3 secondes par page pour l'OCR
        
        ### 🎯 À propos du résumé intelligent:
        
        **Résumé hiérarchique pour longs documents:**
        - 📚 Documents < 1000 mots: Résumé direct
        - 📖 Documents 1000-3000 mots: Résumé optimisé
        - 📕 Documents > 3000 mots: Résumé multi-étapes
          1. Le document est divisé en chunks de 1000 mots
          2. Chaque chunk est résumé séparément
          3. Les résumés sont combinés et re-résumés
          4. Résultat: un résumé cohérent et complet
        
        **Styles de résumé:**
        - 🎯 **Concis**: Points clés essentiels (30-150 mots)
        - ⚖️ **Équilibré**: Vue d'ensemble complète (50-250 mots) - Recommandé
        - 📋 **Détaillé**: Analyse approfondie (80-400 mots)
        
        ### ⚡ Performance:
        - 🟢 **Mode GPU**: Traitement OCR 5-10x plus rapide
        - 🔵 **Mode CPU**: Fonctionne sur tous les systèmes
        - 📊 **Barre de progression**: Suivez chaque étape du traitement
        - 💾 **Export complet**: Sauvegardez tous les résultats en un clic
        """
    )
    
    # Event handlers for PDF
    analyze_pdf_btn.click(
        fn=lambda pdf, min_s, max_s: process_document(pdf, "PDF", min_s, max_s),
        inputs=[pdf_input, min_summary_pdf, max_summary_pdf],
        outputs=[
            extracted_output,
            title_output,
            author_output,
            date_output,
            location_output,
            summary_output,
            page_info_output,
            confidence_output,
            word_count_output,
            page_count_output
        ]
    )
    
    # Event handlers for Images
    analyze_images_btn.click(
        fn=lambda imgs, min_s, max_s: process_document(imgs, "Images", min_s, max_s),
        inputs=[images_input, min_summary_img, max_summary_img],
        outputs=[
            extracted_output,
            title_output,
            author_output,
            date_output,
            location_output,
            summary_output,
            page_info_output,
            confidence_output,
            word_count_output,
            page_count_output
        ]
    )
    
    # Clear button
    clear_btn.click(
        fn=clear_all,
        inputs=None,
        outputs=[
            pdf_input,
            extracted_output,
            title_output,
            author_output,
            date_output,
            location_output,
            summary_output,
            page_info_output,
            confidence_output,
            word_count_output,
            page_count_output
        ]
    )
    
    # Export functionality
    export_btn.click(
        fn=export_results,
        inputs=[
            extracted_output,
            title_output,
            author_output,
            date_output,
            location_output,
            summary_output,
            page_info_output
        ],
        outputs=export_output
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(share=True)