import gradio as gr
import glob
from docx import Document
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer

# 1. Получение текстов из абзацев и таблиц .docx
def get_blocks_from_docx():
    docx_list = glob.glob("*.docx")
    if not docx_list:
        return ["Файл .docx не найден!"]
    doc = Document(docx_list[0])
    blocks = []
    # Абзацы
    for p in doc.paragraphs:
        txt = p.text.strip()
        if txt and not (len(txt) <= 3 and txt.isdigit()):
            blocks.append(txt)
    # Таблицы
    for table in doc.tables:
        for row in table.rows:
            row_text = " | ".join(cell.text.strip() for cell in row.cells if cell.text.strip())
            if row_text:
                blocks.append(row_text)
    # Удаляем дубли
    seen = set()
    uniq_blocks = []
    for b in blocks:
        if b not in seen:
            uniq_blocks.append(b)
            seen.add(b)
    return uniq_blocks

blocks = get_blocks_from_docx()
vectorizer = TfidfVectorizer().fit(blocks)
matrix = vectorizer.transform(blocks)

# 2. Загрузка модели rut5-base-multitask
tokenizer = T5Tokenizer.from_pretrained("cointegrated/rut5-base-multitask")
model = T5ForConditionalGeneration.from_pretrained("cointegrated/rut5-base-multitask")
model.eval()
device = 'cpu'

def rut5_answer(question, context):
    prompt = f"question: {question} context: {context}"
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
    with torch.no_grad():
        output_ids = model.generate(input_ids, max_length=200, num_beams=4)
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

# 3. Комбинированная функция: поиск + генерация
def ask_chatbot(question):
    if not question.strip():
        return "Пожалуйста, введите вопрос."
    if len(blocks) < 2:
        return "Ошибка: база знаний пуста или слишком мала. Проверьте .docx."
    # Находим релевантный абзац
    user_vec = vectorizer.transform([question])
    sims = cosine_similarity(user_vec, matrix)[0]
    best_idx = sims.argmax()
    best_block = blocks[best_idx]
    score = sims[best_idx]
    if score < 0.12:
        context = ""
    else:
        context = best_block
    # Генерируем нейросетевой ответ на русском с учетом найденного контекста
    answer = rut5_answer(question, context)
    # Для большей прозрачности покажем также фрагмент из документа (можно убрать)
    if context:
        return f"**Ответ:** {answer}\n\n---\n**Релевантный фрагмент из документа:**\n{context}"
    else:
        return f"**Ответ:** {answer}\n\n(Контекст в документе не найден — ответ дан на основе общего знания модели.)"

EXAMPLES = [
    "Какие требования к объему магистерской диссертации?",
    "Как оформить список литературы?",
    "Какие сроки сдачи и защиты ВКР?",
    "Что должно быть во введении?",
    "Какой процент оригинальности требуется?",
    "Как оформлять формулы?"
]

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Русскоязычный FAQ-чат-бот на базе вашей методички и нейросетевой модели

        Задайте вопрос — получайте свежий AI-ответ, опирающийся на ваш документ!
        """
    )
    question = gr.Textbox(label="Ваш вопрос", lines=2)
    ask_btn = gr.Button("Получить ответ")
    answer = gr.Markdown(label="Ответ", visible=True)
    ask_btn.click(ask_chatbot, question, answer)
    question.submit(ask_chatbot, question, answer)
    gr.Markdown("#### Примеры вопросов:")
    gr.Examples(EXAMPLES, inputs=question)
    gr.Markdown("""
    ---

    ### Контакты (укажите свои)
    Преподаватель: ___________________
    Email: ___________________________
    Кафедра: _________________________
    """)

demo.launch()