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

pdf_path = "pdf/bialy montaz elektryczny.pdf"

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

print("--- INICJALIZACJA AGENTA ---")
print(f"1. Ładowanie pliku PDF ({pdf_path})...")
loader = PyPDFLoader(pdf_path)
docs = loader.load()

print(f"2. Dzielenie dokumentu ({len(docs)} stron) na fragmenty...")
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
splits = text_splitter.split_documents(docs)

print("3. Tworzenie lokalnej bazy wektorowej (to może chwilę potrwać)...")
embeddings = OllamaEmbeddings(model="qwen2.5-coder:7b")
vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
retriever = vectorstore.as_retriever()

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

print("\nGotowe! Agent jest gotowy do rozmowy.")
print("Wpisz swoje pytanie do dokumentu (lub wpisz 'exit', aby zakończyć).\n")

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

        # Szukamy fragmentów pasujących do pytania
        relevant_docs = retriever.invoke(pytanie)
        kontekst = "\n\n".join([doc.page_content for doc in relevant_docs])

        # Tworzymy prompt dla modelu
        prompt = f"""Odpowiedz na pytanie bazując WYŁĄCZNIE na poniższym fragmencie dokumentu PDF. 
Jeśli nie masz tej informacji w tekście, napisz, że jej nie znasz.

Tekst z dokumentu PDF:
{kontekst}

Pytanie: {pytanie}
"""

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

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