import os
from app_agent_config import AgentConfig
from utils.logger import log_response
from model.custom_agent import CustomHfAgent 
from model.conversation_chain_singleton import ConversationChainSingleton

def cut_text_after_keyword(text, keyword):
    index = text.find(keyword)
    if index != -1:
        return text[:index].strip()
    return text


class Controller:
    def __init__(self):
        self.agent_config = AgentConfig()

    image = []
    def handle_submission(self, user_message ):
        
        log_response("User input \n {}".format(user_message))
        log_response("selected_tools \n {}".format(self.agent_config.selected_tools))
        log_response("url_endpoint \n {}".format(self.agent_config.url_endpoint))
        log_response("document \n {}".format(self.agent_config.document))
        log_response("image \n {}".format(self.agent_config.image))
        log_response("context \n {}".format(self.agent_config.context))
        
        agent = CustomHfAgent(
            url_endpoint=self.agent_config.url_endpoint,
            token=os.environ['HF_token'],
            additional_tools=self.agent_config.elected_tools,
            input_params={"max_new_tokens": 192},
        )

        angent_respone = agent.chat(user_message,document=self.agent_config.document,image=self.agent_config.image, context = self.agent_config.context)

        log_response("Agent Response\n {}".format(angent_respone))

        return angent_respone

    def handle_submission_chat(self, user_message, angent_respone):
        # os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.environ['HF_token']
        agent_chat_bot = ConversationChainSingleton().get_conversation_chain()    

        if angent_respone is not None:
            text = agent_chat_bot.predict(input=user_message + angent_respone)
        else:
            text = agent_chat_bot.predict(input=user_message)

        result = cut_text_after_keyword(text, "Human:")
       
        print(result)

        return result