import os
from time import time
import huggingface_hub
import streamlit as st
from config import config
from functioncall import ModelInference


@st.cache_resource(show_spinner="Loading model..")
def init_llm():
    huggingface_hub.login(token=config.hf_token, new_session=False)
    llm = ModelInference(chat_template=config.chat_template)
    return llm


def get_response(prompt):
    try:
        return llm.generate_function_call(
            prompt, config.chat_template, config.num_fewshot, config.max_depth
        )
    except Exception as e:
        return f"An error occurred: {str(e)}"


def get_output(context, user_input):
    try:
        config.status.update(label=":bulb: Preparing answer..")
        script_dir = os.path.dirname(os.path.abspath(__file__))
        prompt_path = os.path.join(script_dir, 'prompt_assets', 'output_sys_prompt.yml')
        prompt_schema = llm.prompter.read_yaml_file(prompt_path)
        sys_prompt = (
            llm.prompter.format_yaml_prompt(prompt_schema, dict())
            + f"Information:\n{context}"
        )
        convo = [
            {"role": "system", "content": sys_prompt},
            {"role": "user", "content": user_input},
        ]
        response = llm.run_inference(convo)
        return response
    except Exception as e:
        return f"An error occurred: {str(e)}"


def main():
    st.title("LLM-ADE 9B Demo")

    input_text = st.text_area("Enter your text here:", value="", height=200)

    if st.button("Generate"):
        if input_text:
            with st.status("Generating response...") as status:
                config.status = status
                agent_resp = get_response(input_text)
                st.write(get_output(agent_resp, input_text))
                config.status.update(label="Finished!", state="complete", expanded=True)
        else:
            st.warning("Please enter some text to generate a response.")


llm = init_llm()


def main_headless(prompt: str):
    start = time()
    agent_resp = get_response(prompt)
    print("\033[94m" + get_output(agent_resp, prompt) + "\033[0m")
    print(f"Time taken: {time() - start:.2f}s\n" + "-" * 20)


if __name__ == "__main__":
    if config.headless:
        import fire
        fire.Fire(main_headless)
    else:
        main()