from utils.retriever import get_engine
from utils.index import create_index
from utils.constant import INDEX_PATH, DATA_PATH
import os
import shutil


def set_keys(co_api_key: str, llama_cloud_api_key: str) -> str:
    try:
        os.environ["CO_API_KEY"] = co_api_key
        os.environ["LLAMA_CLOUD_API_KEY"] = llama_cloud_api_key
        return "Keys are set successfully"
    
    except Exception as e:
        return str(e)

def handle_file(uploaded_file):
    try:

        temp_file_path = uploaded_file.name
        save_file_path = os.path.join(DATA_PATH, "paper.pdf")
        
        shutil.move(temp_file_path, save_file_path)
        
        create_index(DATA_PATH, INDEX_PATH)
        
        engine_manager = get_engine()[1]
        engine_manager.initialize_index()
        
        return "File uploaded and indexed"
    
    except Exception as e:
        return str(e)

def answer_query(query: str) -> str:
    
    query_engine = get_engine(semantic=True)[0]
    response = query_engine.query(query)
    
    nodes = query_engine.retriever.retrieve(query)
    nodes_str = ""
    
    for node in nodes:

        score = node.get_score()
        text = node.text
        
        nodes_str += f"\nNode: {node.node_id}\nScore: {score:0.3f}\nText: {text[:1000]}\n"
    
    return response, nodes_str

if __name__ == "__main__":
    
    query = "What is Few-NERD"
    response = answer_query(query)
    print(response)