import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_core.messages import HumanMessage
from langchain_community.document_loaders import DirectoryLoader, PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain_community.retrievers import BM25Retriever
import os

folder_path = "pdf"

if not os.path.exists(folder_path):
    print(f"Błąd: Nie znaleziono folderu '{folder_path}'!")
    exit()

print("--- INICJALIZACJA AGENTA HYBRYDOWEGO ---")
print(f"1. Ładowanie dokumentów z folderu '{folder_path}'...")

loader = DirectoryLoader(folder_path, glob="**/*.pdf", loader_cls=PyPDFLoader)
docs = loader.load()

if not docs:
    print("Błąd: Brak plików PDF!")
    exit()

print(f"Załadowano łącznie {len(docs)} stron.")

print("2. Dzielenie dokumentów na fragmenty...")
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
splits = text_splitter.split_documents(docs)

print("3. Tworzenie bazy wektorowej oraz indeksu słów kluczowych...")
embeddings = OllamaEmbeddings(model="qwen2.5-coder:7b")
vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
chroma_retriever = vectorstore.as_retriever(search_kwargs={"k": 5})

# Indeks słów kluczowych (BM25) do dokładnych numerów i indeksów
bm25_retriever = BM25Retriever.from_documents(splits)
bm25_retriever.k = 5

# Własna funkcja łącząca oba podejścia (hybryda)
def hybrid_retrieve(query, k=6):
    docs_chroma = chroma_retriever.invoke(query)
    docs_bm25 = bm25_retriever.invoke(query)
    
    # Łączymy wyniki, usuwając duplikaty po treści
    seen = set()
    combined = []
    for doc in docs_bm25 + docs_chroma:
        if doc.page_content not in seen:
            seen.add(doc.page_content)
            combined.append(doc)
    return combined[:k]

llm = ChatOllama(model="qwen2.5-coder:7b")

print("\nGotowe! Agent hybrydowy czeka na pytania.")
print("Wpisz swoje pytanie (lub 'exit', aby zakończyć).\n")

while True:
    try:
        pytanie = input("Ty: ")
        if not pytanie.strip():
            continue
        if pytanie.lower() in ["exit", "wyjscie", "koniec", "q"]:
            print("Do zobaczenia!")
            break

        # Pobieramy dokumenty hybrydowo
        relevant_docs = hybrid_retrieve(pytanie)
        kontekst = "\n\n".join([doc.page_content for doc in relevant_docs])

        prompt = f"""Odpowiedz na pytanie bazując WYŁĄCZNIE na poniższych fragmentach dokumentów. 
Jeśli nie masz tej informacji w tekstach, napisz, że jej nie znasz.

Teksty źródłowe:
{kontekst}

Pytanie: {pytanie}
"""

        response = llm.invoke([HumanMessage(content=prompt)])
        print(f"\nAgent: {response.content}\n" + "-"*40)

    except KeyboardInterrupt:
        print("\nDo zobaczenia!")
        break