import os
from openai import OpenAI
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

class RAGQueryAnswerAssistant:
    def __init__(self, api_key):
        self.knowledge_base = {}
        self.embeddings = {}
        self.client = OpenAI(api_key = api_key)
        self.model = SentenceTransformer('all-MiniLM-L6-v2')

    def load_file(self, file_path):
        if not file_path.endswith('.txt'):
            print(f"Error: {file_path} is not a .txt file.")
            return

        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
                file_name = os.path.basename(file_path)
                self.knowledge_base[file_name] = content
                self.embeddings[file_name] = self.model.encode([content])[0]
                print(f"Successfully loaded and embedded {file_name}")
        except FileNotFoundError:
            print(f"Error: File {file_path} not found.")
        except Exception as e:
            print(f"Error loading file: {str(e)}")

    def retrieve_relevant_context(self, query, top_k=2):
        query_embedding = self.model.encode([query])[0]
        similarities = {}
        for file_name, embedding in self.embeddings.items():
            similarity = cosine_similarity([query_embedding], [embedding])[0][0]
            similarities[file_name] = similarity
        
        sorted_files = sorted(similarities.items(), key=lambda x: x[1], reverse=True)
        relevant_context = ""
        for file_name, _ in sorted_files[:top_k]:
            relevant_context += f"Content from {file_name}:\n{self.knowledge_base[file_name]}\n\n"
        return relevant_context

    def answer_query(self, query):
        relevant_context = self.retrieve_relevant_context(query)
        prompt = f"""Given the following context:

{relevant_context}

Please answer the following question:
{query}

If the answer is not explicitly stated in the context, use your knowledge to provide the most relevant and accurate answer possible. If you cannot provide an answer based on the given context or your knowledge, please state that clearly."""

        try:
            response = self.client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant that answers questions based on the given context."},
                    {"role": "user", "content": prompt}
                ]
            )
            return response.choices[0].message['content']
        except Exception as e:
            return f"Error generating response: {str(e)}"

    def list_loaded_files(self):
        if self.knowledge_base:
            return "Loaded files:\n" + "\n".join(self.knowledge_base.keys())
        else:
            return "No files have been loaded yet."

# Example usage
# Clave de API de OpenAI
api_key = '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'
assistant = RAGQueryAnswerAssistant(api_key)

# Main loop
while True:
    print("\nGPT Query Answer Assistant")
    print("1. Load a file")
    print("2. Ask a question")
    print("3. List loaded files")
    print("4. Exit")
    
    choice = input("Enter your choice (1-4): ")
    
    if choice == '1':
        file_path = input("Enter the path to the .txt file: ")
        assistant.load_file(file_path)
    elif choice == '2':
        query = input("Enter your question: ")
        answer = assistant.answer_query(query)
        print(answer)
    elif choice == '3':
        print(assistant.list_loaded_files())
    elif choice == '4':
        print("Goodbye!")
        break
    else:
        print("Invalid choice. Please try again.")