import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import faiss
import numpy as np
import ollama
import gradio as gr
import os
import argparse
from sentence_transformers import SentenceTransformer

# === KONFIGURACJA ===
model_name = "hse.ably.do:latest"  # Nazwa modelu Ollama
faiss_index_path = "faiss_index.idx"  # Plik indeksu FAISS
kodeks_file = "./docs/kodekspracy.txt"  # Plik z treścią kodeksu pracy
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")  # Model do embedowania tekstu

# === KROK 1: WCZYTYWANIE KODEKSU PRACY ===
def load_kodeks(filepath):
    with open(filepath, 'r', encoding='utf-8') as file:
        content = file.read()
        articles = content.split('\n\n')
        documents = []
        for article in articles:
            if article.strip().startswith('Art.'):
                documents.append(article.strip())
    return documents

# === KROK 2: TWORZENIE INDEKSU FAISS ===
def create_faiss_index(sections):
    embeddings = embedding_model.encode(sections, convert_to_numpy=True)  # Tworzenie wektorów
    index = faiss.IndexFlatL2(embeddings.shape[1])  # Indeks FAISS
    index.add(embeddings)  # Dodanie wektorów do FAISS
    faiss.write_index(index, faiss_index_path)  # Zapis indeksu
    return index, sections

# === KROK 3: WYSZUKIWANIE NAJBLIŻSZEGO FRAGMENTU ===
def search_faiss(query, index, sections):
    query_vector = embedding_model.encode([query], convert_to_numpy=True)  # Wektor zapytania
    _, idx = index.search(query_vector, 1)  # Szukamy 1 najbliższego sąsiada
    return sections[idx[0][0]]  # Zwracamy najbardziej pasujący fragment kodeksu

# === KROK 4: GENEROWANIE ODPOWIEDZI Z OLLAMA ===
def generate_response(user_query):
    if not os.path.exists(faiss_index_path):
        return "Błąd: Indeks FAISS nie istnieje. Uruchom aplikację z opcją --rebuild-index, aby go wygenerować."

    try:
        index = faiss.read_index(faiss_index_path)  # Wczytanie indeksu FAISS
    except Exception as e:
        return f"Błąd: Nie można załadować indeksu FAISS. Spróbuj odbudować go za pomocą --rebuild-index. Szczegóły: {str(e)}"

    sections = load_kodeks(kodeks_file)  # Wczytanie kodeksu
    best_match = search_faiss(user_query, index, sections)  # Znalezienie najbliższego fragmentu

    # Kontekst dla Ollama
    prompt = f"Odpowiedz na pytanie na podstawie następującego tekstu:\n\n{best_match}\n\nPytanie: {user_query}"

    response = ollama.chat(model=model_name, messages=[{"role": "user", "content": prompt}])
    return response.get("message", "Błąd: Nie udało się wygenerować odpowiedzi.")

# === KROK 5: URUCHOMIENIE INTERFEJSU WEBOWEGO ===
iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(label="Zadaj pytanie o kodeks pracy"),
    outputs=gr.Textbox(label="Odpowiedź"),
    title="Asystent Kodeksu Pracy",
    description="Wpisz pytanie, a system zwróci odpowiedni fragment kodeksu pracy."
)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--rebuild-index", action="store_true", help="Wymuś odbudowanie indeksu FAISS")
    args = parser.parse_args()

    if args.rebuild_index or not os.path.exists(faiss_index_path):
        print("Tworzenie nowego indeksu FAISS...")
        sections = load_kodeks(kodeks_file)
        create_faiss_index(sections)
    else:
        print("Indeks FAISS już istnieje.")
    
    iface.launch(share=True)