import requests
import json
import os
import gradio as gr
import spaces

token = os.getenv('HUGGINGFACEHUB_API_TOKEN')
API_URL = "https://kp4xdy196cw81uf3.us-east-1.aws.endpoints.huggingface.cloud"
headers = {
	"Accept" : "application/json",
	"Authorization": "Bearer " + token,
	"Content-Type": "application/json"
}

def query(payload):
	response = requests.post(API_URL, headers=headers, json=payload)
	return response.json()
    
@spaces.GPU
def generate_response(text):
  input = {
    "inputs": text,
    "parameters": {
      "max_new_tokens" : 128,
      "top_k": 10,
      "top_p": 0.95,
      "typical_p": 0.95,
      "temperature": 0.01,
      "repetition_penalty": 1.03,
      "stop" : ["/nHuman:", "/nUser:", "<end of message>\n"]
    }
  }
  
  output = query(input)
  response = output[0]["generated_text"]
  response = response.replace(text, "")
  splitResponse = response.split("Human:")
  return splitResponse[0]


def format_chat_prompt(message, chat_history):
    prompt = """ Do not repeat questions and do not generate answer for user/human.
    
    You are a helpful hotel booking asssitant.
    Below is an instruction that describes a task.
    Write a response that appropriately completes the request.
    Reply with the most helpful and logic answer. During the conversation you need to ask the user
    the following questions to complete the hotel booking task.
    1) Where would you like to stay and when?
    2) How many people are staying in the room?
    3) Do you prefer any ammenities like breakfast included or gym?
    4) What is your name, your email address and phone number?
    When name, email and phone number is provided, you, a hotel booking assistant, shouldd respond with a thankful answer.
    When the user says please book the room or Yes, book it, you should repsond with : "Yes,I booked the rooom."
    """
    for turn in chat_history:
        user_message, bot_message = turn
        prompt = f"{prompt}\nHuman: {user_message}\nAssistant: {bot_message}"
    prompt = f"{prompt}\nHuman: {message}\nAssistant:"
    return prompt

def chat_output(message, history):
    print(message)
    prompt = format_chat_prompt(message,history)
    result = generate_response(text = prompt)
    return result

with gr.Blocks() as demo:

    #chatbot_component = gr.Chatbot(height=300, label = "history")
    #textbox_component = gr.Textbox(placeholder="Can I help you to book a hotel?", container=False, label = "input", scale=7)

    demo.chatbot_interface = gr.ChatInterface(
        fn=chat_output,
        examples = ["Hello I would like to book a hotel room.", "Hello I want to stay in Nuremberg in 30th of May until 1st of June." , "2 people, me and my partner."
                    "2 people, me and my partner. I would like a breakfast included.", " I would like a breakfast included.", "Just me.I am travelling alone."],
        #outputs=chatbot_component,
        title = "Hotel Booking Assistant Chat 🤗",
        description = "I am your hotel booking assistant. Feel free to start chatting with me."
    )
demo.launch(debug=True)