import gradio as gr 
from transformers import pipeline, AutoTokenizer 

from os import getenv
from huggingface_hub import login 
from history import get_history, update_history

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

tokenizer = AutoTokenizer.from_pretrained("google/gemma-1.1-2b-it")

pipe = pipeline("text-generation", model="google/gemma-1.1-2b-it", revision="refs/pr/14",trust_remote_code=True) # trust here

def talk(prompt,history=[]):  
  
  # Combine the history into a single string
  combined_history = " ".join([entry["content"] for entry in history])
  # Tokenize and count tokens
  tokens = tokenizer.encode(combined_history)

  print('tokens',len(tokens))
 
  #If no history, get history from database
  group_name=getenv("DEFAULT_GROUP_NAME")
  if not len(history):
    old_history=get_history(group_name)
    history=old_history

  print('the history',history)

  # Transform to array of tuples
  transformed_data = []
  user_message = None

  for item in history:
    if item['role'] == 'user':
        user_message = item['content']  # Save user message temporarily
    elif item['role'] == 'assistant' and user_message:
        transformed_data.append((user_message, item['content']))
        user_message = None  # Reset after pairing

  # Output result
  print("transformed_data",transformed_data)

  response=''

  try: 
    for partial_response in pipe(prompt, history=transformed_data, stream=True, max_new_tokens=128000):
        response = partial_response
        yield partial_response
  except Exception as e:
    print("Stream queue was empty. No response generated.",str(e))
    yield "Error: No response generated." 

  #update the history in database
  if response:
    transformed_data.append((prompt,response))
    update_history(group_name,transformed_data)


def initialize():
  messages=[]
  history=get_history(getenv("DEFAULT_GROUP_NAME"))
  for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})
  return messages
 

demo = gr.ChatInterface(
  fn=talk, 
  type="messages",
  chatbot=gr.Chatbot(value=initialize(),type="messages"), 
  examples=[["hi",[]]],
  cache_examples=False,
  title="Streaming")

demo.launch(share=True,ssr_mode=False)