import gradio as gr
from transformers import pipeline
import pdfplumber
import os

# --- Ленивая загрузка модели ---
summarizer = None

def load_model():
    global summarizer
    if summarizer is None:
        summarizer = pipeline(
            "summarization",
            model="sshleifer/distilbart-cnn-12-6"
        )
    return summarizer

# --- Безопасное чтение PDF ---
def extract_pdf_text(path):
    text = ""
    try:
        with pdfplumber.open(path) as pdf:
            for page in pdf.pages:
                chunk = page.extract_text()
                if chunk:
                    text += chunk + "\n"
    except Exception as e:
        return "", f"Ошибка при чтении PDF: {e}"
    return text, None

# --- Универсальное чтение файла ---
def read_file(path):
    if not path:
        return "", "Файл не передан."
    path = str(path).strip()

    if path.lower().endswith(".pdf"):
        return extract_pdf_text(path)

    try:
        with open(path, "r", encoding="utf-8", errors="ignore") as f:
            return f.read(), None
    except Exception as e:
        return "", f"Ошибка при чтении TXT: {e}"

# --- Основная функция ---
def summarize_file(path):
    text, err = read_file(path)
    if err:
        return f"⚠️ {err}"

    if not text.strip():
        return "⚠️ Не удалось извлечь текст."

    if len(text) < 80:
        return "⚠️ Слишком мало текста для суммаризации."

    model = load_model()

    # Ограничиваем текст для избежания OOM
    text = text[:4000]

    try:
        summary = model(
            text,
            max_length=180,
            min_length=60,
            do_sample=False
        )
        return summary[0]["summary_text"]
    except Exception as e:
        return f"⚠️ Ошибка суммаризации: {e}"

# --- Интерфейс Gradio ---
demo = gr.Interface(
    fn=summarize_file,
    inputs=gr.File(type="filepath", label="Загрузите файл (.pdf или .txt)"),
    outputs=gr.Textbox(label="Результат суммаризации"),
    title="Eroha Summarizer 🧠",
    description="Загрузите документ (PDF или TXT), и модель создаст краткое резюме."
)

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