import os
import sys
import string
import pickle # Библиотека для сохранения фрагментов текста на диск

# Загрузчики и разбивка текста
from langchain_community.document_loaders import DirectoryLoader, PyMuPDFLoader, Docx2txtLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

# Эмбеддинги и векторная база
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS

# Ретриверы
from langchain_community.retrievers import BM25Retriever
from langchain_classic.retrievers import EnsembleRetriever

# Цепочки
from langchain_classic.chains.combine_documents import create_stuff_documents_chain
from langchain_classic.chains import create_retrieval_chain

# Базовые модули и LLM
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

# ==========================================
# НАСТРОЙКА ПУТЕЙ
# ==========================================
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
BASE_REGULATIONS_PATH = os.path.join(BASE_DIR, "search_bot", "Регламенты", "АОХКСибцем", "07 Служба Вице-президента по экономике и финансам", "ДИТ")
MODEL_PATH = os.path.join(BASE_DIR, "models", "mpnet_model")

# Пути для сохранения базы данных
FAISS_DB_PATH = os.path.join(BASE_DIR, "search_bot", "faiss_index")
CHUNKS_PATH = os.path.join(BASE_DIR, "search_bot", "chunks.pkl")

# ==========================================
# 1. ИНИЦИАЛИЗАЦИЯ ЭМБЕДДИНГОВ
# ==========================================
print("Загрузка модели эмбеддингов (оффлайн)...")
embedding = HuggingFaceEmbeddings(
    model_name=MODEL_PATH,
    model_kwargs={'device': 'cpu'},
    encode_kwargs={'normalize_embeddings': False}
)

# ==========================================
# 2. ЗАГРУЗКА ИЛИ СОЗДАНИЕ БАЗЫ ДАННЫХ
# ==========================================
# Проверяем, существует ли уже сохраненная база
if os.path.exists(FAISS_DB_PATH) and os.path.exists(CHUNKS_PATH):
    print("✅ Найдена сохраненная база данных! Загружаю с диска (это быстро)...")
    
    # Загружаем фрагменты текста для BM25
    with open(CHUNKS_PATH, 'rb') as f:
        split_docs = pickle.load(f)
        
    # Загружаем векторы FAISS
    vector_store = FAISS.load_local(
        FAISS_DB_PATH, 
        embedding, 
        allow_dangerous_deserialization=True # Обязательно для локальных файлов
    )
else:
    print("⚠️ Сохраненная база не найдена. Начинаю чтение и индексацию документов...\n")
    
    if not os.path.exists(BASE_REGULATIONS_PATH):
        print("❌ ОШИБКА: Указанный путь к документам не существует.")
        sys.exit(1)

    # Загружаем PDF и DOCX
    pdf_loader = DirectoryLoader(BASE_REGULATIONS_PATH, glob="**/*.pdf", loader_cls=PyMuPDFLoader, show_progress=True)
    docx_loader = DirectoryLoader(BASE_REGULATIONS_PATH, glob="**/*.docx", loader_cls=Docx2txtLoader, show_progress=True)
    
    docs = pdf_loader.load() + docx_loader.load()
    if len(docs) == 0:
        print("❌ В указанной папке нет файлов .pdf или .docx. Завершение.")
        sys.exit(1)

    print(f"✅ Загружено страниц/файлов: {len(docs)}. Разбивка на фрагменты...")
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1200, chunk_overlap=200, separators=["\n\n", "\n", ".", " "])
    split_docs = text_splitter.split_documents(docs)

    print("Создаю векторную базу FAISS...")
    vector_store = FAISS.from_documents(split_docs, embedding=embedding)
    
    # СОХРАНЕНИЕ НА ДИСК ДЛЯ БУДУЩИХ ЗАПУСКОВ
    print("💾 Сохраняю базу данных на диск...")
    vector_store.save_local(FAISS_DB_PATH)
    with open(CHUNKS_PATH, 'wb') as f:
        pickle.dump(split_docs, f)
    print("✅ База успешно сохранена!")

# ==========================================
# 3. НАСТРОЙКА РЕТРИВЕРОВ (ПОИСКА)
# ==========================================
# 3.1 Семантический поиск
embedding_retriever = vector_store.as_retriever(search_kwargs={"k": 4})

# 3.2 Лексический поиск
def tokenize(s):
    return s.lower().translate(str.maketrans("", "", string.punctuation)).split(" ")

bm25_retriever = BM25Retriever.from_documents(
    documents=split_docs,
    preprocess_func=tokenize,
    k=5
)

# 3.3 Гибридный поиск (Ансамбль)
ensemble_retriever = EnsembleRetriever(
    retrievers=[embedding_retriever, bm25_retriever],
    weights=[0.4, 0.6]
)

# ==========================================
# 4. ПОДКЛЮЧЕНИЕ К LLAMA.CPP
# ==========================================
print("Подключение к серверу llama.cpp...")
llm = ChatOpenAI(
    base_url="http://127.0.0.1:8080/v1", 
    api_key="not-needed",                
    temperature=0.0,                     
    max_tokens=1024                      
)

prompt = ChatPromptTemplate.from_template('''Ты — строгий корпоративный ИИ-помощник. 
Твоя задача — отвечать на вопросы строго на основании текста предоставленных регламентов.
Внимательно изучи контекст. Если там есть ответ, сформулируй его четко и по делу. Обязательно указывай номер пункта или название документа, откуда взята информация.
Если в контексте НЕТ ответа на вопрос, не придумывай информацию, а выведи фразу: "В предоставленных регламентах нет информации по данному вопросу."

Контекст: 
{context}

Вопрос пользователя: {input}

Ответ:'''
)

document_chain = create_stuff_documents_chain(llm=llm, prompt=prompt)
rag_chain = create_retrieval_chain(ensemble_retriever, document_chain)

# ==========================================
# 5. ТЕСТИРОВАНИЕ СИСТЕМЫ
# ==========================================
print("\n🚀 Система готова к работе! Выполняем запросы...\n")

# Тестовые вопросы по Положению № ПОЛ-177
questions = [
    "Где работник обязан хранить электронные документы, связанные с производственной деятельностью?",
    "Через какой срок удаляются электронные документы, если их не открывали, и можно ли их восстановить?",
    "Кому разрешен доступ к сетевым ресурсам Soft/Distr?",
    "Какие правила хранения и удаления установлены для медиафайлов в сетевой папке 'СВК'?",
    "Разрешено ли самостоятельно предоставлять общий доступ к папкам на своем рабочем компьютере?",
    "Что имеют право сделать сотрудники ДИТ или СИБ, если файл несет угрозу ресурсам?"
]

for q in questions:
    print(f"❓ Вопрос: {q}")
    try:
        response = rag_chain.invoke({'input': q})
        print(f"🤖 Ответ: {response['answer']}\n")
        
        # Распечатка источников (полезно для отладки)
        print("🔍 Найденные источники:")
        for i, doc in enumerate(response['context']):
             source_name = doc.metadata.get('source', 'Неизвестный источник').split('/')[-1] # Берем только имя файла
             print(f"   [{i+1}] Файл: {source_name}")
        print("\n" + "-" * 50 + "\n")
        
    except Exception as e:
        print(f"❌ Ошибка при запросе к серверу llama.cpp: {e}")
        print("-" * 50)