import time
import random
from tool_helper import tool_list, parse_and_execute_tool_call
from tool_functions import register_dummy
from inference import Inference, torch_reseed
import datetime



messages = []
inference = None

# systemmessage at the very begin of the chat. Will be concatenated with the automatic tool usage descriptions
systemmessage = "Hold a casual conversation with the user. Keep responses short at max 3 sentences. Answer using markdown to the user."

# system message for role flip so the model automatically answers for the user
roleflip = {"role": "system", "content": "Keep the conversation going, ask for more information on the subject. Keep messages short at max 1-2 sentences. Do not thank and say goodbye."}

# system messages and user message to bring the model to summarize the entire conversation
summarize = {"role": "system", "content": "Summarize the conversation as a single, cohesive paragraph. Avoid using any bullet points, numbers, or list formatting. Write in plain text with natural sentences that flow together seamlessly."}
summarize_user = {"role": "system", "content": "Can you summarize the conversation?"}

# system message to create a conversation title
title_prompt = {"role": "system", "content": "Please create a very short and descriptive title or label for this conversation. Maximum 2-5 words. Use only plain text, avoid numbering, special characters, or unnecessary formatting-focus on clarity and brevity."}

append_toolcalls = True

register_dummy()




def append_generate_chat(input_text: str, role="user"):
    t_start = time.time()

    # generate AI response
    if input_text != None:
        messages.append({"role": role, "content": input_text})

    inputs = inference.tokenize(messages, tokenize=True)

    outputs, out_text = inference.generate(inputs)

    # append result to message history
    messages.append({"role": "assistant", "content": out_text})

    print("")
    print("generation took %.3fs (%d tokens)" % (time.time() - t_start, len(outputs[0])))

    # handle tool call and check if a tool call has happened.
    tool_result = parse_and_execute_tool_call(out_text, tool_list)
    if tool_result != None:
        # tool call happened
        tool_result = "<tool_response>%s</tool_response>" % tool_result
        # depending on the chat template the tool response tags must or must not be passed. :(
        append_generate_chat(tool_result, role="tool")




def main():
    global messages
    global inference

    inference = Inference()

    current_date_and_time = datetime.datetime.now().strftime("Current date is %Y-%m-%d and its %H:%M %p right now.")
    if append_toolcalls:
        messages = [{"role": "system", "content": systemmessage + "\n" + current_date_and_time + "\n" + inference.generate_tool_use_header(tool_list)}]
    else:
        messages = [{"role": "system", "content": systemmessage + "\n" + current_date_and_time}]

    while True:
        # print an input prompt to receive text or commands
        input_text = input(">>> ")
        print("")


        if input_text.startswith("!"):
            append_generate_chat("<tool_response>%s</tool_response>" % input_text[1:], role="tool")
            # append_generate_chat("%s" % input_text[1:], role="tool")  # depending on the chat template the tool response tags must or must not be passed. :(

        elif input_text.startswith("/clear"):
            print("clearing chat history")
            start_msg = messages[0]
            messages = [start_msg]
            print("")

        elif input_text.startswith("/history"):
            history = inference.tokenize(messages, tokenize=False)
            # history = tokenizer.apply_chat_template(messages, return_tensors="pt", tokenize=False, add_generation_prompt=False)
            print(history)

        elif input_text.startswith("/undo"):
            if len(messages) > 2:
                print("undo latest prompt")
                messages = messages[:-2]
            else:
                print("cannot undo because there are not enough messages on history.")
            print("")

        elif input_text.startswith("/regen"):
            if len(messages) >= 2:
                print("regenerating message (not working)")
                messages = messages[:-1]
                seed = random.randint(0, 2**32 - 1)  # Generate a random seed
                torch_reseed(seed)
                append_generate_chat(None)
            else:
                print("cannot regenerate because there are not enough messages on history.")
            print("")

        elif input_text.startswith("/more"):
            append_generate_chat(None)

        elif input_text.startswith("/file"):
            filename = input_text[len("/file "):]
            print("read '%s' for prompt:" % filename)
            with open(filename, "r") as f:
                content = f.read()
            print(content)
            append_generate_chat(content)

        elif input_text.startswith("/auto"):
            messages_backup = messages
            messages = [roleflip]
            for m in messages_backup:
                role = m["role"]
                content = m["content"]
                if role == "user":
                    role = "assistant"
                elif role == "assistant":
                    role = "user"
                if role != "system":
                    messages.append({"role": role, "content": content})
            append_generate_chat(None)  # will automatically advance the conversation as 'user'
            last_message = messages[-1]
            last_message["role"] = "user"
            messages = messages_backup + [last_message]
            append_generate_chat(None)  # 'regular' chatbot answer

        elif input_text.startswith("/summarize"):
            messages_temp = list(filter(lambda x: x["role"] != "system", messages))
            messages_temp = [summarize] + messages_temp + [summarize_user]  # copy dict in last instance
            # messages_temp[-1]["role"] = "user"
            input_ids = inference.tokenize(messages_temp, tokenize=True, assistant_prefix="The conversation was about ")
            generated_tokens, full_output = inference.generate(input_ids)

        elif input_text.startswith("/title"):
            messages_temp = list(filter(lambda x: x["role"] != "system", messages))
            messages_temp = [title_prompt] + messages_temp #+ [dict(title)]  # copy dict in last instance
            messages_temp[-1]["role"] = "user"
            input_ids = inference.tokenize(messages_temp, tokenize=True, assistant_prefix="Title: ")
            generated_tokens, full_output = inference.generate(input_ids)

        elif input_text.startswith("/help"):
            print("!<prompt>   answer as 'tool' in <tool_response> tags")
            print("/clear      clear chat history")
            print("/undo       undo latest prompt")
            print("/regen      regenerate the last message")
            print("/more       generate more additional information")
            print("/file       read prompt input from file")
            print("/auto       automatically advance conversation")
            print("/summarize  generate a summary of the chat")
            print("/title      generate a title of the chat")
            print("/help       print this message")
            print("")

        elif input_text.startswith("/"):
            print("unknown command.")

        else:
            append_generate_chat(input_text)



if __name__ == "__main__":
    main()