import gradio as gr
import os
from langdetect import detect
from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering
import numpy as np
import re
import random

# Загрузка и предварительная обработка текстовых файлов
def load_and_preprocess_files():
    files = {
        "vampires": "vampires.txt",
        "werewolves": "werewolves.txt",
        "humans": "humans.txt"
    }
    
    knowledge_base = {}
    for category, filename in files.items():
        try:
            with open(filename, 'r', encoding='utf-8') as file:
                content = file.read()
                # Разбиваем на осмысленные блоки (абзацы)
                paragraphs = [p.strip() for p in content.split('\n\n') if p.strip()]
                knowledge_base[category] = paragraphs
        except FileNotFoundError:
            print(f"Файл {filename} не найден")
            knowledge_base[category] = []
    
    return knowledge_base

# Инициализация модели вопрос-ответ
def initialize_qa_model():
    tokenizer = AutoTokenizer.from_pretrained('DeepPavlov/rubert-base-cased')
    model = AutoModelForQuestionAnswering.from_pretrained('DeepPavlov/rubert-base-cased')
    qa_pipeline = pipeline('question-answering', model=model, tokenizer=tokenizer)
    return qa_pipeline

# Поиск релевантной информации
def find_relevant_context(question, knowledge_base):
    all_paragraphs = []
    for category, paragraphs in knowledge_base.items():
        all_paragraphs.extend(paragraphs)
    
    # Чтобы не работать по всей базе, берём случайные 10 абзацев (упрощённый вариант, можно сделать лучше)
    sample_paragraphs = random.sample(all_paragraphs, min(10, len(all_paragraphs)))
    context = " ".join(sample_paragraphs)
    return context

# Генерация ответа через модель
def generate_answer(question, context, qa_pipeline):
    result = qa_pipeline(question=question, context=context)
    return result['answer']

# Обработка вопроса
def process_question(question, history):
    try:
        if detect(question) != 'ru':
            return "Пожалуйста, задавайте вопросы на русском языке.", history
    except:
        pass
    
    if not hasattr(process_question, 'knowledge_base'):
        process_question.knowledge_base = load_and_preprocess_files()
    
    if not hasattr(process_question, 'qa_pipeline'):
        process_question.qa_pipeline = initialize_qa_model()
    
    context = find_relevant_context(question, process_question.knowledge_base)
    answer = generate_answer(question, context, process_question.qa_pipeline)
    
    if not answer.strip():
        answer = "Извините, я не смог найти точный ответ. Попробуйте переформулировать вопрос."

    history.append((question, answer))
    return "", history

# Создание интерфейса
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""<h1 style='text-align: center'>🧛‍♂️ Мир сверхъестественного 🐺</h1>""")
    gr.Markdown("""<div style='text-align: center'>Задавайте вопросы о вампирах, оборотнях и людях на русском языке</div>""")
    
    msg = gr.Textbox(
        label="Ваш вопрос",
        placeholder="Введите вопрос и нажмите Enter...",
        container=False
    )
    
    examples = gr.Examples(
        examples=[
            "Какие слабости у вампиров?",
            "Как защититься от оборотней?",
            "Чем люди отличаются от других существ?",
            "Расскажи подробнее о вампирах"
        ],
        inputs=[msg],
        label="Примеры вопросов:"
    )
    
    chatbot = gr.Chatbot(
        label="Диалог",
        height=500
    )
    
    with gr.Row():
        submit = gr.Button("Отправить", variant="primary")
        clear = gr.Button("Очистить историю")
    
    submit.click(process_question, [msg, chatbot], [msg, chatbot])
    msg.submit(process_question, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

demo.launch()