import tkinter as tk
from tkinter import filedialog
import sys
import ollama

def print_progress(iteration, total, prefix='', suffix='', decimals=1, bar_length=50):
    """Afișează o bară de progres în linia de comandă."""
    percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
    filled_length = int(bar_length * iteration // total)
    bar = '#' * filled_length + '-' * (bar_length - filled_length)
    sys.stdout.write(f'\r{prefix} |{bar}| {percent}% {suffix}')
    sys.stdout.flush()
    if iteration == total:
        print()

def load_text_file():
    """Folosește o fereastră grafică doar pentru a selecta fișierul TXT și procesează textul."""
    root = tk.Tk()
    root.withdraw()  # Ascundem fereastra principală

    file_path = filedialog.askopenfilename(title="Selectează un fișier TXT", filetypes=[("Text files", "*.txt")])
    if not file_path:
        print("Nu a fost selectat niciun fișier!")
        exit()

    try:
        with open(file_path, "r", encoding="utf-8") as file:
            text_content = file.read()
    except Exception as e:
        print(f"Eroare la încărcarea fișierului: {e}")
        exit()

    # Împărțim textul în fragmente folosind dublu newline ca delimitator.
    chunks = [chunk.strip() for chunk in text_content.split("\n\n") if chunk.strip()]

    # Afișăm lungimea fiecărui fragment (numărul de caractere) pentru informație.
    print("Lungimea fragmentelor (în caractere):")
    for i, chunk in enumerate(chunks, start=1):
        print(f"Fragment {i}: {len(chunk)} caractere")

    # Calculăm embedding-urile pentru fiecare fragment și afișăm progresul în consolă.
    total = len(chunks)
    embeddings = []
    print("\nSe procesează fișierul:")
    for i, chunk in enumerate(chunks, start=1):
        print_progress(i, total, prefix="Progres", suffix="Finalizat", bar_length=50)
        emb = ollama.embeddings(model="mxbai-embed-large", prompt=chunk)["embedding"]
        embeddings.append(emb)

    print("Fișierul a fost încărcat și procesat cu succes!")
    return chunks, embeddings

def cosine_similarity(vec1, vec2):
    """Calculează similaritatea cosinus între doi vectori."""
    dot_product = sum(a * b for a, b in zip(vec1, vec2))
    norm1 = sum(a * a for a in vec1) ** 0.5
    norm2 = sum(b * b for b in vec2) ** 0.5
    return dot_product / (norm1 * norm2) if norm1 and norm2 else 0

def generate_natural_answer(question, context):
    """
    Generează un răspuns natural folosind modelul llama3.2:latest.
    Modelul primește întrebarea și contextul extins.
    """
    prompt = f"Întrebare: {question}\nContext: {context}\nRăspunde natural, în limba română, bazându-te pe context:"
    response = ollama.generate(model="llama3.2:latest", prompt=prompt)
    return response["response"]

if __name__ == "__main__":
    # Încarcă fișierul folosind fereastra de dialog (fără a afișa interfața principală)
    chunks, embeddings = load_text_file()

    # Loop de întrebări/răspunsuri în linia de comandă
    while True:
        question = input("\nIntrodu întrebarea (sau tastează 'exit' pentru a încheia): ").strip()
        if not question or question.lower() in ("exit", "quit", "q"):
            print("Programul se încheie.")
            break

        # Calculăm embedding-ul pentru întrebare
        question_embedding = ollama.embeddings(model="mxbai-embed-large", prompt=question)["embedding"]

        # Calculăm similaritatea între întrebare și fiecare fragment din fișier
        similarities = [cosine_similarity(question_embedding, emb) for emb in embeddings]

        # Selectăm primele 12 fragmente cu cea mai mare similaritate pentru context extins
        top_n = 12
        top_indices = sorted(range(len(similarities)), key=lambda i: similarities[i], reverse=True)[:top_n]
        context = "\n\n".join([chunks[i] for i in top_indices])

        # Generează răspunsul pe baza contextului extins folosind modelul llama3.2:latest
        natural_answer = generate_natural_answer(question, context)

        # Afișează răspunsul și indică fragmentele folosite (numerotate de la 1)
        used_fragments = sorted([i + 1 for i in top_indices])
        print("\nRăspuns:")
        print(natural_answer)
        print(f"\n(Context: fragmentele {used_fragments})")

        # Întreabă dacă se dorește vizualizarea contextului complet
        show_ctx = input("\nDoriți să vizualizați contextul complet? (da/nu): ").strip().lower()
        if show_ctx in ("da", "d"):
            print("\nContext complet folosit:")
            print(context)
            # Întreabă dacă se dorește salvarea contextului într-un fișier TXT
            save_ctx = input("\nDoriți să salvați contextul complet într-un fișier TXT? (da/nu): ").strip().lower()
            if save_ctx in ("da", "d"):
                file_name = input("Introduceți numele fișierului (ex: context.txt): ").strip()
                if not file_name:
                    file_name = "context.txt"
                try:
                    with open(file_name, "w", encoding="utf-8") as f:
                        f.write(context)
                    print(f"Contextul a fost salvat în fișierul {file_name}.")
                except Exception as e:
                    print(f"Eroare la salvarea fișierului: {e}")
        else:
            print("\nUrmătoarea întrebare:")