"""
Chat engine.

TODOs:
- Better prompts.
- Output reader / parser.
- Agents for evaluation and task planning / splitting.
    * Haystack for orchestration
- Tools for agents
    * Haystack for orchestration
- 

"""


client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")

def query_submit(user_message, history):
    return "", history + [[user_message, None]]
  
def format_prompt(query, history, lookback):
    prompt = "Responses should be no more than 100 words long.\n"
    
    for previous_query, prevous_completion in history[-lookback:]:
        prompt += f"<s>[INST] {previous_query} [/INST] {prevous_completion}</s> "
    
    prompt += f"[INST] {query} [/INST]"
  
    return prompt

def query_completion(
    query,
    history,
    lookback = 3,
    max_new_tokens = 256,
):

    generateKwargs = dict(
        max_new_tokens = max_new_tokens,
        seed = 1337,
    )

    formatted_query = format_prompt(query, history, lookback)
    
    stream = client.text_generation(
        formatted_query,
        **generateKwargs,
        stream = True,
        details = True,
        return_full_text = False
    )
    
    history[-1][1] = ""
    
    for response in stream:
        history[-1][1] += response.token.text
        yield history

        
"""
Chat UI using Gradio Blocks.

Blocks preferred for lower-level "atomic" layout control and state management.

TODOs:
- State management for dynamic components update.
- Add scratpad readout to right of chat log.
    * Placeholder added for now.
- Add functionality to retry button.
    * Placeholder added for now.
- Add dropdown for model selection.
- Add textbox for HF model selection.
    
"""

with gr.Blocks() as chatUI:
    with gr.Row():
        chatOutput = gr.Chatbot(
            bubble_full_width = False,
            scale = 2
        )
        agentWhiteBoard = gr.Markdown(scale = 1)
            
    with gr.Row():
        queryInput = gr.Textbox(
            placeholder = "Please enter you question or request here...",
            show_label = False,
            scale = 4
        )
        submitButton = gr.Button("Submit", scale = 1)
        
    with gr.Row():
        retry = gr.Button("Retry (null)")
        clear = gr.ClearButton([queryInput, chatOutput])
    
    # gr.State()
    
    queryInput.submit(
        fn = query_submit,
        inputs = [queryInput, chatOutput],
        outputs = [queryInput, chatOutput],
        queue = False,
    ).then(
        fn = query_completion,
        inputs = [queryInput, chatOutput],
        outputs = [chatOutput],
    )
    
    submitButton.click(
        fn = query_submit,
        inputs = [queryInput, chatOutput],
        outputs = [queryInput, chatOutput],
        queue = False,
    ).then(
        fn = query_completion,
        inputs = [queryInput, chatOutput],
        outputs = [chatOutput],
    )

chatUI.queue()
chatUI.launch(show_api = False)