import gradio as gr
import easyocr
import numpy as np
from PIL import Image
from transformers import pipeline
from docx import Document
from pptx import Presentation
from spellchecker import SpellChecker
from langdetect import detect
import fitz  # PyMuPDF
import os
import tempfile

# -------- OCR --------
reader = easyocr.Reader(["fr", "en", "ar"], gpu=False)

# -------- Résumés par langue --------
summarizers = {
    "fr": pipeline("summarization", model="facebook/bart-large-cnn"),
    "en": pipeline("summarization", model="facebook/bart-large-cnn"),
    "ar": pipeline("summarization", model="csebuetnlp/mT5_multilingual_XLSum")
}

# -------- Correcteur OCR --------
spell = SpellChecker(language="fr")

def correct_ocr(text):
    words = text.split()
    corrected = []
    for w in words:
        if w.isalpha():
            corrected.append(spell.correction(w) or w)
        else:
            corrected.append(w)
    return " ".join(corrected)

# -------- Résumé --------
def summarize_text(text, lang):
    wc = len(text.split())
    if wc < 30:
        return "Texte trop court pour être résumé."

    model = summarizers.get(lang, summarizers["fr"])
    result = model(
        text,
        max_length=min(150, wc),
        min_length=min(30, wc // 2),
        do_sample=False
    )
    return result[0]["summary_text"]

# -------- IMAGE --------
def process_image(image, lang):
    image = np.array(image)
    results = reader.readtext(image, paragraph=True)
    text = " ".join([r[1] for r in results])

    if not text.strip():
        return "Aucun texte détecté.", "", None

    corrected = correct_ocr(text)
    summary = summarize_text(corrected, lang)

    file_path = save_summary(summary)
    return corrected, summary, file_path

# -------- PDF --------
def extract_text_pdf(path):
    doc = fitz.open(path)
    return "\n".join(page.get_text() for page in doc)

# -------- FICHIER --------
def extract_text(file_path):
    ext = os.path.splitext(file_path)[1].lower()

    if ext == ".txt":
        return open(file_path, encoding="utf-8", errors="ignore").read()

    if ext == ".docx":
        return "\n".join(p.text for p in Document(file_path).paragraphs)

    if ext == ".pptx":
        prs = Presentation(file_path)
        return "\n".join(
            shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text")
        )

    if ext == ".pdf":
        return extract_text_pdf(file_path)

    return ""

def process_file(file, lang):
    text = extract_text(file.name)

    if not text.strip():
        return "Aucun texte détecté.", "", None

    corrected = correct_ocr(text)
    summary = summarize_text(corrected, lang)

    file_path = save_summary(summary)
    return corrected, summary, file_path

# -------- Télécharger --------
def save_summary(summary):
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".txt")
    tmp.write(summary.encode("utf-8"))
    tmp.close()
    return tmp.name

# -------- INTERFACE --------
with gr.Blocks(title="OCR & Résumé Intelligent") as demo:
    gr.Markdown("""
    # 🧠 OCR & Résumé Intelligent
    ✔ Image | TXT | DOCX | PPTX | PDF  
    ✔ Résumé FR / AR / EN  
    ✔ Correction OCR  
    ✔ Téléchargement du résumé
    """)

    lang = gr.Radio(
        choices=[("Français", "fr"), ("العربية", "ar"), ("English", "en")],
        value="fr",
        label="Langue du résumé"
    )

    with gr.Tabs():

        with gr.Tab("🖼️ Image"):
            img = gr.Image(type="pil")
            btn_img = gr.Button("Extraire & Résumer", variant="primary")
            txt_img = gr.Textbox(label="Texte corrigé", lines=10)
            sum_img = gr.Textbox(label="Résumé", lines=5)
            dl_img = gr.File(label="📥 Télécharger le résumé")

            btn_img.click(
                process_image,
                inputs=[img, lang],
                outputs=[txt_img, sum_img, dl_img]
            )

        with gr.Tab("📄 Fichier"):
            file = gr.File(file_types=[".txt", ".docx", ".pptx", ".pdf"])
            btn_file = gr.Button("Extraire & Résumer", variant="primary")
            txt_file = gr.Textbox(label="Texte corrigé", lines=10)
            sum_file = gr.Textbox(label="Résumé", lines=5)
            dl_file = gr.File(label="📥 Télécharger le résumé")

            btn_file.click(
                process_file,
                inputs=[file, lang],
                outputs=[txt_file, sum_file, dl_file]
            )

if __name__ == "__main__":
    demo.launch()