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
   


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

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

        response = agent.chat(user_message,document=document,image=image, context = context)

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

        return response

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


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

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

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

        return result