import gradio as gr
from huggingface_hub import InferenceClient 
from typing import List, Union, Dict, Tuple

from os import getenv
from huggingface_hub import login 

from history import get_history, update_history

# Login to Hugging Face
login(getenv("Token")) 

#initialize the inference client using the name of the 
#inference endpoint: meta-llama/Llama-3.2-3B-Instruct
client = InferenceClient("meta-llama/Llama-3.2-3B-Instruct")


#Transform gradio history by breaking any tuple into 2 dicts
def transform_gradio_history(
        history: List[Union[Dict[str, str], 
        Tuple[str, str]]]
        ) -> List[Dict[str, str]]:
    transformed_history = []
    
    for entry in history:
        if (isinstance(entry, list) or isinstance(entry, tuple)) and len(entry) == 2:
            transformed_history.append({"role": "user", "content": entry[0]})
            transformed_history.append({"role": "assistant", "content": entry[1]})
        elif isinstance(entry, dict):
            transformed_history.append(entry)
    
    return transformed_history

def respond(
    message,
    history: list[dict[str, str]],
    system_message, #system prompt
    max_tokens,
    temperature,
    top_p, 
    group_name #user Id
):
    if not group_name:
        #user must pass user Id to the group_name.
        #This is used to identify the user
        yield "User ID required"
    else:
        messages=history
           
        #If no history, get history from database 
        if not len(messages):
            messages=get_history(group_name) 
        
        #Break any tuples into 2 dicts
        messages=transform_gradio_history(messages)
        
        #Add prompt to list of messages
        messages.append({"role": "user", "content": message}) 

        response = ""

        #Create new list of all messages, starting with system prompt
        mainMessage=[{"role": "system", "content": system_message}, *messages] 

        for msg in client.chat_completion(
            mainMessage, #message list
            max_tokens=max_tokens,
            stream=True, #enable streaming of response
            temperature=temperature,
            top_p=top_p, 
        ):
            token = msg.choices[0].delta.content

            response += token
            yield response

        #update the conversation history in database
        if response:
            messages.append({"role": "assistant", "content": response})
            update_history(group_name,messages)

def initialize():
  messages=[] 
  return messages

demo = gr.ChatInterface(
    respond,
    type="messages",
    chatbot=gr.Chatbot(value=initialize(),type="messages"),
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),  
        gr.Textbox( label="User ID"),
    ],
)


if __name__ == "__main__":
    demo.launch(share=True,ssr_mode=False)