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

import numpy as np
import faiss
import gradio as gr
import torch
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM
import ollama

# 1️⃣ Inicjalizacja modelu do embeddingów
embed_model = SentenceTransformer("all-MiniLM-L6-v2")

# 2️⃣ Dodanie dokumentów i embeddingów
def read_documents_from_file(file_path):
    with open(file_path, '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
#documents = [
#    "Jak założyć firmę w Polsce?", 
#    "Jak rozliczyć podatek VAT?", 
#    "Procedura składania reklamacji w e-sklepie.",
#    "Jakie dokumenty są potrzebne do rejestracji działalności?"
#]
file_path = './docs/kodekspracy.txt'  # Zmień na właściwą ścieżkę
documents = read_documents_from_file(file_path)

# 3️⃣ Wygenerowanie embeddingów
embeddings = embed_model.encode(documents)

# 4️⃣ Inicjalizacja FAISS
dim = embeddings.shape[1]  # Wymiar embeddingu
index = faiss.IndexFlatL2(dim)
index.add(np.array(embeddings, dtype=np.float32))

# 5️⃣ Wczytanie modelu Ollama (Gemma 2)
model_name = "hse.ably.do:latest"  # Ścieżka do modelu w systemie
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b")

# 6️⃣ Funkcja wyszukiwania w FAISS
def search(query, k=5):
    query_embedding = embed_model.encode([query])  # Przekształć zapytanie w embedding
    _, indices = index.search(np.array(query_embedding, dtype=np.float32), k)  # Szukaj w indeksie FAISS
    return indices  # Zwróć indeksy najbardziej podobnych dokumentów

# 7️⃣ Funkcja generowania odpowiedzi z Ollama
def generate_response(query):
    indices = search(query)  # Znajdź najbardziej podobne dokumenty
    relevant_documents = [documents[i] for i in indices[0]]  # Pobierz dokumenty na podstawie wyników wyszukiwania
    
    prompt = " ".join(relevant_documents) + " " + query  # Przygotuj prompt
    response = ollama.chat(model=model_name, messages=[{"role": "user", "content": prompt}])  # Generowanie odpowiedzi przez Ollama
    
    return response["text"]

# 8️⃣ Interfejs Gradio (Open-WebUI)
iface = gr.Interface(fn=generate_response, inputs="text", outputs="text")

iface.launch(share=True)  # Uruchom interfejs