import os
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain_community.llms import HuggingFaceHub
import gradio as gr
import re

# 1. Загрузка и очистка всех .txt файлов
def load_documents(folder_path):
    documents = []
    for file_name in os.listdir(folder_path):
        if file_name.endswith(".txt"):
            loader = TextLoader(os.path.join(folder_path, file_name), encoding="utf-8")
            docs = loader.load()
            for doc in docs:
                # Очищаем спецсимволы типа [=/ и прочую ерунду
                doc.page_content = re.sub(r'\[=/.*?\]', '', doc.page_content)
                documents.append(doc)
    return documents

# 2. Разбивка на чанки
def split_documents(documents):
    splitter = RecursiveCharacterTextSplitter(chunk_size=700, chunk_overlap=100)
    return splitter.split_documents(documents)

# 3. Создание эмбеддингов
def create_embeddings():
    return HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

# 4. Загрузка модели
def load_llm():
    return HuggingFaceHub(
        repo_id="IlyaGusev/saiga_mistral_7b_gguf",  # можно заменить на что-то другое, если будет падать
        model_kwargs={"temperature": 0.6, "max_new_tokens": 300}
    )

# 5. Построение цепочки
def build_qa_chain():
    raw_docs = load_documents("lore")  # Папка lore/ рядом с app.py
    docs = split_documents(raw_docs)
    embeddings = create_embeddings()
    db = FAISS.from_documents(docs, embeddings)
    retriever = db.as_retriever()
    llm = load_llm()
    return RetrievalQA.from_chain_type(llm=llm, retriever=retriever)

# 6. Интерфейс
qa_chain = build_qa_chain()

def answer_question(question):
    result = qa_chain.run(question)
    return result

iface = gr.Interface(fn=answer_question, inputs="text", outputs="text", title="Чат по Лору (RU)")
iface.launch()