import gradio as gr
from langchain_community.document_loaders import UnstructuredMarkdownLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain.vectorstores import FAISS
from langchain_community.llms import HuggingFaceHub
from langchain.prompts import ChatPromptTemplate
from dotenv import load_dotenv
import os
import shutil

# Загрузка переменных окружения
load_dotenv()

CHROMA_PATH = "/tmp/chroma"
DATA_PATH = ""  # Укажите путь к вашим данным, например "data", если файл не в корне
PROMPT_TEMPLATE = """
Ответь на вопрос, используя только следующий контекст:
{context}
---
Ответь на вопрос на основе приведенного контекста: {question}
"""

# Глобальная переменная для статуса
status_message = "Инициализация..."

def initialize_chroma():
    global status_message
    status_message = "Инициализация векторной базы..."
    else:
        status_message = "База данных Chroma уже существует. Пересоздаем базу данных..."
        shutil.rmtree(CHROMA_PATH)
        
    generate_data_store()
    status_message = "База данных Chroma создана и подготовлена."
    
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
        cache_folder="/tmp/model_cache",
        model_kwargs={'device': 'cpu'},
        encode_kwargs={'normalize_embeddings': True}
    )
    db = Chroma(
        persist_directory=CHROMA_PATH,
        embedding_function=embeddings
    )
    return db

def generate_data_store():
    documents = load_documents()
    if not documents:
        raise Exception("Документы не загружены!")
    chunks = split_text(documents)
    status_message = f"Первые 50 символов первого чанка: {chunks[0].page_content[:50]}"
    save_to_chroma(chunks)
    
    
def load_documents():
    file_path = os.path.join(DATA_PATH, "pl250320252.md")
    if not os.path.exists(file_path):
        global status_message
        status_message = f"Ошибка: Файл {file_path} не найден."
        return []
    loader = UnstructuredMarkdownLoader(file_path)
    documents = loader.load()
    return documents

def split_text(documents: list[Document]):
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=900,
        chunk_overlap=300,
        length_function=len,
        add_start_index=True,
    )
    chunks = text_splitter.split_documents(documents)
    global status_message
    status_message += f"\nРазделено {len(documents)} документов на {len(chunks)} частей."
    return chunks

def save_to_chroma(chunks: list[Document]):
    if os.path.exists(CHROMA_PATH):
        shutil.rmtree(CHROMA_PATH)
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
        model_kwargs={'device': 'cpu', 'trust_remote_code': True},
        encode_kwargs={'normalize_embeddings': False}  # Экспериментируйте с этим
    )
    Chroma.from_documents(
        chunks,
        embeddings,
        persist_directory=CHROMA_PATH
    )

def process_query(query_text: str, db):
    results = db.similarity_search_with_relevance_scores(query_text, k=3)
    global status_message
    status_message += f"\nНайдено {len(results)} результатов с релевантностью: {[round(score, 2) for _, score in results]}"
    
    if not results:
        return "Не найдено подходящих результатов.", []
    
    if results[0][1] < 0.5:
        response = "Нет высококачественных совпадений. Возможные варианты:\n"
        response += "\n".join([f"- {doc.page_content[:50]}..." for doc, _ in results])
        return response, [doc.metadata.get("source") for doc, _ in results]
    
    context_text = "\n\n---\n\n".join([doc.page_content for doc, _ in results])
    prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
    prompt = prompt_template.format(context=context_text, question=query_text)
    model = HuggingFaceHub(
        repo_id="google/flan-t5-small",
        model_kwargs={"temperature": 0.5, "max_length": 512}
    )
    response_text = model.predict(prompt)
    sources = [doc.metadata.get("source", None) for doc, _ in results]
    return response_text, sources

def chat_interface(query_text):
    global status_message
    db = initialize_chroma()
    response, sources = process_query(query_text, db)
    full_response = f"{status_message}\n\nОтвет: {response}\n\nИсточники: {', '.join(sources) if sources else 'Нет источников'}"
    return full_response

# Создание папок
os.makedirs("/tmp/model_cache", exist_ok=True)
os.makedirs("/tmp/chroma", exist_ok=True)

# Интерфейс Gradio
interface = gr.Interface(
    fn=chat_interface,
    inputs=gr.Textbox(lines=2, placeholder="Введите ваш вопрос здесь..."),
    outputs="text",
    title="Чат с документами",
    description="Задайте вопрос, и я отвечу на основе загруженных документов."
)

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