from sortinghat import SortingHat
import random
import torch
import gradio as gr
from PIL import Image

from peft import AutoPeftModelForCausalLM

from transformers import AutoTokenizer
from time import sleep

class HPChatBot:
    def __init__(self,
                 model_path: str = 'rnltls/harrypotter_lexicon_finetune',
                 device_map: str = 'auto',
                 load_in_4_bit: bool = True,
                 **quant_kwargs) -> None:
        self.model = None
        self.tokenizer = None
        self.image_processor = None
        self.conv = None
        self.conv_img = []
        self.img_tensor = []
        self.roles = None
        self.stop_key = None
        self.is_chat = False
        self.is_waldo = False
        self.load_models(model_path,
                         device_map=device_map,
                         load_in_4_bit=load_in_4_bit,
                         **quant_kwargs)

    def load_models(self, model_path: str,
                    device_map: str,
                    load_in_4_bit: bool,
                    **quant_kwargs) -> None:
        model = AutoPeftModelForCausalLM.from_pretrained(
            model_path, # YOUR MODEL YOU USED FOR TRAINING
            load_in_4bit = load_in_4_bit,
        )
        tokenizer = AutoTokenizer.from_pretrained(model_path)

        self.model = AutoPeftModelForCausalLM.from_pretrained(
            "rnltls/harrypotter_lexicon_finetune", # YOUR MODEL YOU USED FOR TRAINING
            load_in_4bit = load_in_4_bit,
        )
        self.tokenizer = AutoTokenizer.from_pretrained(model_path,
                                                       use_fast=False)
    
    def generate_answer(self, prompt):
        output = self.model.generate(**prompt, max_new_tokens = 256)
        generated_text = self.tokenizer.decode(output[0], skip_special_tokens=True)
        return generated_text

HogwartGemma = "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/HogwartGemma.jpg?raw=true"
SpellProf = ["https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Spell_stand.jpg?raw=true",
           "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Spell_think.jpg?raw=true",
           "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Spell_show.jpg?raw=true"]
PotionProf = ["https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Potion_stand.jpg?raw=true",
              "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Potion_think.jpg?raw=true",
              "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Potion_show.jpg?raw=true"]
OtherProf = ["https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Other_stand.jpg?raw=true",
             "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Other_think.jpg?raw=true",
             "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Other_show.jpg?raw=true"]
HouseHat = ["https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/sorting_hat.jpg?raw=true",
          "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/sorting_hat_think.jpg?raw=true",
          "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/sorting_hat_gryffindor.png?raw=true",
            "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/sorting_hat_hufflepuff.png?raw=true",
            "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/sorting_hat_ravenclaw.png?raw=true",
          "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/sorting_hat_slytherin.png?raw=true"]
Icons = ["https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/wand_icon.png?raw=true",
         "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/cauldron.png?raw=true",
         "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/golden-snitch.png?raw=true",
         "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/hat_icon.png?raw=true"]
enable_btn = gr.Button(interactive=True)
disable_btn = gr.Button(interactive=False)
spell_chat, potion_chat, other_chat, house_chat = False, False, False, False
question_counter = 0
quit_counter = 0

question_list = None
quote_list = None
answer_list = []
question_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{}

### Input:
{}

### Response:
{}"""

def show_image(path):
    # Convert To PIL Image
    image = Image.open(path)
    return image

def spell_chatting(chat_history, request: gr.Request):
    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, answer_list
    answer_list = []
    quit_counter=0
    question_counter = 0
    chat_history.clear()
    spell_chat, potion_chat, other_chat, house_chat = True, False, False, False
    bot_message = "Welcome to Magic Spell Class!\nTell me what you want to achieve, and I’ll suggest the perfect spell for it!\nIf you ask in the format: 'What spell can I use when I ~?', I can give you even better suggestions!"
    txt_box = gr.Textbox(value="What spell can I use when I ", interactive=True)
    prof_IMG = SpellProf[0]
    chat_history.append([None, bot_message])
    return chat_history, prof_IMG, txt_box, enable_btn, enable_btn, enable_btn

def potion_chatting(chat_history, request: gr.Request):
    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, answer_list
    answer_list = []
    quit_counter=0
    question_counter = 0
    chat_history.clear()
    spell_chat, potion_chat, other_chat, house_chat = False, True, False, False
    bot_message = "Welcome to Potion Class!\nTell me what you want to achieve, and I’ll suggest the perfect potion for it!\nIf you ask in the format: 'What potion can I make when I ~?', I can give you even better suggestions!"
    txt_box = gr.Textbox(value="What potion can I make when I ", interactive=True)
    prof_IMG = PotionProf[0]
    chat_history.append([None,bot_message])
    return chat_history, prof_IMG, txt_box, enable_btn, enable_btn, enable_btn

def other_chatting(chat_history, request: gr.Request):
    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, answer_list
    answer_list = []
    quit_counter=0
    question_counter = 0
    chat_history.clear()
    spell_chat, potion_chat, other_chat, house_chat = False, False, True, False
    bot_message = "Welcome to Library!\nWelcome to the library! Feel free to ask me anything if you're curious!"
    txt_box = gr.Textbox(placeholder="Ask anything you're curious about in the Wizarding World!", value="", interactive=True)
    prof_IMG = OtherProf[0]
    chat_history.append([None,bot_message])
    return chat_history, prof_IMG, txt_box, enable_btn, enable_btn, enable_btn

def house_chatting(chat_history, request: gr.Request):
    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, question_counter, question_list, quote_list, answer_list
    answer_list = []
    question_list = sortinghat.choose_questions()
    quote_list = sortinghat.choose_quotes()
    quit_counter=0
    question_counter = 0
    chat_history.clear()
    spell_chat, potion_chat, other_chat, house_chat = False, False, False, True
    bot_message = "Hmm, let's see… where shall I place you?\n" + question_list[question_counter]
    txt_box = gr.Textbox(placeholder="Answer the question", value="", interactive=True)
    prof_IMG = HouseHat[0]
    chat_history.append([None,bot_message])
    return chat_history, prof_IMG, txt_box, enable_btn, enable_btn, enable_btn

def ask_question(chat_history, text_data, request: gr.Request):
    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, answer_list
    if spell_chat:
        print(text_data)
        prof_IMG = SpellProf[1]
        chat_history.append([text_data, None])
        text_data = gr.Textbox(value="What spell can I use when I ", interactive=True)
        return chat_history, prof_IMG, text_data

    if potion_chat:
        print(text_data)
        prof_IMG = PotionProf[1]
        chat_history.append([text_data, None])
        text_data = gr.Textbox(value="What potion can I make when I ", interactive=True)
        return chat_history, prof_IMG, text_data

    if other_chat:
        print(text_data)
        prof_IMG = OtherProf[1]
        chat_history.append([text_data, None])
        text_data = gr.Textbox(placeholder="Ask anything you're curious about in the Wizarding World!", value="", interactive=True)
        return chat_history, prof_IMG, text_data

    if house_chat:
        print(text_data)
        prof_IMG = HouseHat[1]
        chat_history.append([text_data, None])
        answer_list.append(text_data)
        text_data = gr.Textbox(placeholder="Answer the question", value = "", interactive=True)
        return chat_history, prof_IMG, text_data

def clean_chatting(chat_history, request: gr.Request):
    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, question_counter, question_list, quote_list, answer_list
    quit_counter = 0
    question_counter = 0
    question_list=sortinghat.choose_questions()
    quote_list=sortinghat.choose_quotes()
    answer_list=[]
    chat_history.clear()
    if spell_chat:
        bot_message = "Welcome to Magic Spell Class!\nTell me what you want to achieve, and I’ll suggest the perfect spell for it!\nIf you ask in the format: 'What spell can I use when I ~?', I can give you even better suggestions!"
        prof_IMG = SpellProf[0]
    if potion_chat:
        bot_message = "Welcome to Potion Class!\nTell me what you want to achieve, and I’ll suggest the perfect potion for it!\nIf you ask in the format: 'What potion can I make when I ~?', I can give you even better suggestions!"
        prof_IMG = PotionProf[0]
    if other_chat:
        bot_message = "Welcome to Library!\nWelcome to the library! Feel free to ask me anything if you're curious!"
        prof_IMG = OtherProf[0]
    if house_chat:
        bot_message = "Hmm, let's see… where shall I place you?\n" + question_list[question_counter]
        prof_IMG = HouseHat[0]
    chat_history.append([None, bot_message])
    return chat_history, prof_IMG

def run_model(chat_history, request: gr.Request):
    global spell_chat, potion_chat, other_chat, house_chat, question_counter, question_list, quote_list, answer_list

    if question_counter < 4 and house_chat:
        bot_message = quote_list[question_counter] + '\n' + question_list[question_counter+1]
        question_counter += 1
        chat_history[-1][1] = bot_message
        return chat_history

    elif house_chat:
        print("Choosing the house")
        print(question_list)
        print(answer_list)
        user_responses = "\n".join([f"{i+1}. {q}\nAnswer: {a}" for i, (q,a) in enumerate(zip(question_list, answer_list))])

        prompt = f"""
            You are the Sorting Hat at Hogwarts. Based on the student's answers to the following questions, assign them to the most suitable house (Gryffindor, Hufflepuff, Ravenclaw, or Slytherin) and provide three reasons for your decision.
            Gryffindor values courage, nerve, and chivalry.
            Hufflepuff values hard work, patience, justice, and loyalty.
            Ravenclaw values intelligence, learning, wisdom, and wit.
            Slytherin values ambition, cunning, leadership, and resourcefulness.

            ### Instruction:
            Please format the response like this:
            'Your house is [house]!
            Here's why:
            1. [reason 1]
            2. [reason 2]
            3. [reason 3]
            Welcome to your new house at Hogwarts!'

            ### Input:
            {user_responses}

            ### Response:
        """
        inputs = HP_chatbot.tokenizer([prompt], return_tensors="pt").to("cuda")
        output = HP_chatbot.generate_answer(inputs)
        output = output.split("Response:")[-1]
        print(output)
        chat_history[-1][1] = output
        return chat_history

    alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

    ### Instruction:
    {}

    ### Input:
    {}

    ### Response:
    {}"""

    inputs = HP_chatbot.tokenizer(
    [
        alpaca_prompt.format(
            chat_history[-1][0], # instruction
            "", # input
            "", # output - leave this blank for generation!
        )
    ], return_tensors = "pt").to("cuda")
    output = HP_chatbot.generate_answer(inputs)

    output = output.split("Response:")[-1]
    chat_history[-1][1] = output
    return chat_history

def end_chatting(chat_history, textbox, request: gr.Request):
    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, question_counter
    if not house_chat or quit_counter > 0:
        quit_counter = 0
        question_counter = 0
        chat_history.clear()
        chat_history.append([None, "Welcome to Hogwarts!"])
        textbox = gr.Textbox(show_label=False, placeholder="Welcome to Hogwart!", value= "", container=False, interactive = False)
        return chat_history, HogwartGemma, textbox, disable_btn, disable_btn, disable_btn
    else:
        quit_counter += 1
        chat_history.append([None, "Quit, you say? Oh, I see your plight, But the housing question’s still in sight.\
                             \nThough weary you may be, don't yet retreat,For there's still a task you must complete!"])
        return chat_history, HouseHat[0], textbox, enable_btn, enable_btn, enable_btn

def result_img(chat_history, textbox, submit_btn,  request: gr.Request):
    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, question_counter
    if spell_chat:
        return SpellProf[2], textbox, submit_btn
    if potion_chat:
        return PotionProf[2], textbox, submit_btn
    if other_chat:
        return OtherProf[2], textbox, submit_btn
    if house_chat:
        if question_counter == 4:
            quit_counter = 1
            house = chat_history[-1][1]
            # print(house)
            if "gryffindor" in house.lower():
                insignia =  HouseHat[2]
            elif "hufflepuff" in house.lower():
                insignia =  HouseHat[3]
            elif "ravenclaw" in house.lower():
                insignia =  HouseHat[4]
            elif "slytherin" in house.lower():
                insignia =  HouseHat[5]
            else:
                return HouseHat[0], textbox, disable_btn
            textbox = gr.Textbox(show_label=False, placeholder="Housing Done.", value= "", container=False, interactive = False)
            return insignia, textbox, disable_btn
        else:
            return HouseHat[0], textbox, submit_btn

def build_gradio(concurrency_count=10):
    textbox = gr.Textbox(show_label=False, placeholder="Welcome to Hogwart!", container=False, interactive = False)
    with gr.Blocks(
            theme='gstaff/xkcd',
        ) as demo:
        state = gr.State()
        with gr.Row():
            with gr.Column(scale=3):

                imagebox = gr.Image(value=HogwartGemma, type="pil", interactive=False, show_label=False, show_download_button=False, show_fullscreen_button=False)

                spell_btn = gr.Button(icon=Icons[0], value="Spell", interactive = True, scale = 2)
                potion_btn = gr.Button(icon=Icons[1], value="Potion", interactive = True, scale = 2)
                other_btn = gr.Button(icon=Icons[2], value="Others", interactive = True, scale = 2)
                house_btn = gr.Button(icon=Icons[3], value="Find your House", interactive = True, scale = 2)

            with gr.Column(scale=8):
                initial_message = [None,"Welcome to Hogwarts!"]
                chatbot = gr.Chatbot(
                    label='Free to ask!',
                    value= [initial_message],
                    height=600,
                )
                with gr.Row():
                    with gr.Column(scale=8):
                        textbox.render()
                    with gr.Column(scale=1, min_width=50):
                        submit_btn = gr.Button(value="Send", variant="primary", interactive = False)

                with gr.Row(elem_id="buttons") as button_row:
                    finish_btn = gr.Button(value="🏁 End Talking", interactive = False)
                    clear_btn = gr.Button(value="🗑️  Clear", interactive=False)

        box_list = [spell_btn, potion_btn, other_btn, house_btn, textbox, submit_btn, finish_btn, clear_btn]

        spell_btn.click(
            spell_chatting,
            inputs=[chatbot],
            outputs=[chatbot, imagebox, textbox, submit_btn, finish_btn, clear_btn]
        )
        potion_btn.click(
            potion_chatting,
            inputs=[chatbot],
            outputs=[chatbot, imagebox, textbox, submit_btn, finish_btn, clear_btn]
        )
        other_btn.click(
            other_chatting,
            inputs=[chatbot],
            outputs=[chatbot, imagebox, textbox, submit_btn, finish_btn, clear_btn]
        )
        house_btn.click(
            house_chatting,
            inputs=[chatbot],
            outputs=[chatbot, imagebox, textbox, submit_btn, finish_btn, clear_btn]
        )
        textbox.submit(
            ask_question,
            inputs = [chatbot, textbox],
            outputs = [chatbot, imagebox, textbox]
        ).then(
            run_model,
            inputs = [chatbot],
            outputs = [chatbot]
        ).then(
            result_img,
            inputs=[chatbot, textbox, submit_btn],
            outputs=[imagebox, textbox, submit_btn]
        )

        submit_btn.click(
            ask_question,
            inputs = [chatbot, textbox],
            outputs = [chatbot, imagebox, textbox]
        ).then(
            run_model,
            inputs = [chatbot],
            outputs = [chatbot]
        ).then(
            result_img,
            inputs=[chatbot, textbox, submit_btn],
            outputs=[imagebox, textbox, submit_btn]
        )

        finish_btn.click(
            end_chatting,
            inputs=[chatbot, textbox],
            outputs=[chatbot, imagebox, textbox, submit_btn, finish_btn, clear_btn]
        )

        clear_btn.click(
            clean_chatting,
            inputs=[chatbot],
            outputs=[chatbot, imagebox]
        )

    return demo

import argparse
if __name__ == "__main__":

    # 컨트롤러에서 사용가능한 모델 가져옴. 나는 모델 하나만 쓸 것이기 때문에, get_model_list()에서 해당 모델이 동작하고 있는 url을 넘겨주면 된다!
    # models = [args.model_url]

    # HP_chatbot = HPChatBot(load_in_8bit=True,
    #                    bnb_8bit_compute_dtype=torch.float16,
    #                    bnb_8bit_use_double_quant=True,
    #                    bnb_8bit_quant_type='nf8')

    sortinghat = SortingHat()

    # Gradio를 이용해서 데모 만들기
    demo = build_gradio()
    demo.queue(
        api_open=False
    ).launch(
        server_port=7861,
        share=True
    )