"""
 /*************************************************************************
 * 
 * CONFIDENTIAL
 * __________________
 * 
 *  Copyright (2023-2024) AI Labs, IronOne Technologies, LLC
 *  All Rights Reserved
 *
 *  Author  : Theekshana Samaradiwakara
 *  Description :Python Backend API to chat with private data  
 *  CreatedDate : 14/11/2023
 *  LastModifiedDate : 18/03/2024
 *************************************************************************/
 """

import os
import time
import logging
logger = logging.getLogger(__name__)
from dotenv import load_dotenv
from reggpt.chains.llmChain import get_qa_chain
from reggpt.output_parsers.output_parser import qa_chain_output_parser

from reggpt.configs.model import QA_MODEL_TYPE
from reggpt.utils.retriever import load_ensemble_retriever
load_dotenv()

verbose = os.environ.get('VERBOSE')

qa_model_type=QA_MODEL_TYPE
# retriever=load_faiss_retriever()
retriever=load_ensemble_retriever()
# retriever=load_multi_query_retriever(multi_query_model_type)
logger.info("retriever loaded:")

qa_chain= get_qa_chain(qa_model_type,retriever)

def run_qa_chain(query):
    try:
        logger.info(f"run_qa_chain : Question: {query}")
        # Get the answer from the chain
        start = time.time()
        # res = qa_chain(query)
        res = qa_chain.invoke({"question": query, "chat_history":""})
        if 'I dont know' in res["answer"] or "don't know" in res["answer"]:
            res['answer'] = "I currently do not have the information to answer this question. Please rephrase your question or ask another question." 
        # res = response
        # answer, docs = res['result'],res['source_documents']
        end = time.time()

        # log the result
        logger.info(f"Answer (took {round(end - start, 2)} s.) \n: {res}")

        return qa_chain_output_parser(res)
    
    except Exception as e:
        logger.exception(e)
        raise e

def run_qa_chain_answer_only(query):
    try:
        logger.info(f"run_qa_chain : Question: {query}")
        # Get the answer from the chain
        start = time.time()
        # res = qa_chain(query)
        res = qa_chain.invoke({"question": query, "chat_history":""})
        # res = response
        # answer, docs = res['result'],res['source_documents']
        end = time.time()

        # log the result
        logger.info(f"Answer (took {round(end - start, 2)} s.) \n: {res}")

        return res["answer"]
    
    except Exception as e:
        logger.exception(e)
        raise e