import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import pdfplumber

# ---- TEK MODEL: Türkçe özet + QA için mT5 ----
MODEL_NAME = "mukayese/mt5-base-turkish-summarization"

device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)


# ---- Yardımcı fonksiyonlar ----

def extract_pdf_text(pdf_file) -> str:
    """PDF dosyasından düz metin çıkar."""
    if pdf_file is None:
        return ""

    text_pages = []
    with pdfplumber.open(pdf_file.name) as pdf:
        for page in pdf.pages:
            page_text = page.extract_text() or ""
            text_pages.append(page_text)

    full_text = "\n\n".join(text_pages).strip()
    return full_text


def generate_text(prompt: str, max_new_tokens: int = 256) -> str:
    """mT5 ile verilen prompt'a göre metin üret."""
    if not prompt:
        return "Metin boş görünüyor."

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=1024,
    ).to(device)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_length=max_new_tokens,
            num_beams=4,
            early_stopping=True,
        )

    text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return text.strip()


def summarize_pdf(pdf_file):
    text = extract_pdf_text(pdf_file)
    if not text:
        return "PDF'ten metin çıkarılamadı. Dosya boş veya okunamıyor olabilir."

    prompt = (
        "Aşağıdaki Türkçe metni kısa ve anlaşılır bir şekilde özetle.\n\n"
        f"Metin:\n{text}\n\n"
        "Özet:"
    )

    return generate_text(prompt, max_new_tokens=256)


def qa_on_pdf(pdf_file, question):
    text = extract_pdf_text(pdf_file)
    if not text:
        return "Önce geçerli bir PDF yüklemelisin."
    if not question:
        return "Lütfen PDF hakkında bir soru yaz."

    prompt = (
        "Aşağıdaki Türkçe metne göre soruya cevap ver. Bilmediğin şeyi uydurma.\n\n"
        f"Metin:\n{text}\n\n"
        f"Soru: {question}\n"
        "Cevap:"
    )

    return generate_text(prompt, max_new_tokens=256)


# ---- Gradio arayüzü ----

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🧠 ZenkaMind PDF Analiz (hafif sürüm)

        - PDF yükle
        - Özet al
        - Aynı PDF hakkında soru sor
        (Şu an tek model: mT5 Türkçe özet/QA)
        """
    )

    with gr.Row():
        pdf_input = gr.File(label="PDF dosyası yükle", file_types=[".pdf"])

    question = gr.Textbox(
        label="PDF hakkında soru sor (opsiyonel)",
        placeholder="Örnek: Bu PDF'in ana fikri ne?",
    )

    with gr.Row():
        btn_summary = gr.Button("📄 PDF'yi Özetle")
        btn_qa = gr.Button("❓ Soruyu Cevapla")

    summary_output = gr.Textbox(
        label="Özet",
        lines=10,
    )

    answer_output = gr.Textbox(
        label="Cevap",
        lines=10,
    )

    btn_summary.click(
        fn=summarize_pdf,
        inputs=[pdf_input],
        outputs=[summary_output],
    )

    btn_qa.click(
        fn=qa_on_pdf,
        inputs=[pdf_input, question],
        outputs=[answer_output],
    )


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