import gradio as gr
from huggingface_hub import InferenceClient
import random

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
#client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
#client = InferenceClient("DrDomedag/LocoLama")
#client = InferenceClient("unsloth/llama-3.2-1b-instruct-bnb-4bit")
#client = InferenceClient("unsloth/llama-3.2-1b-bnb-4bit")
#client = InferenceClient("unsloth/Llama-3.2-1B-Instruct") # Works!
#client = InferenceClient("unsloth/Llama-3.2-3B-Instruct") # Works!
#client = InferenceClient("DrDomedag/LocoLamav2")
#client = InferenceClient("T3lli/lora_model")



words = [
    # Nouns (250)
    "apple", "bridge", "cat", "door", "engine", "forest", "giraffe", "horizon", "island", "jungle",
    "kite", "lake", "mountain", "notebook", "ocean", "penguin", "quartz", "rainbow", "snowflake", "tornado",
    "umbrella", "village", "waterfall", "xylophone", "yard", "zebra", "actor", "ball", "camera", "desert",
    "elephant", "firefly", "garden", "hat", "igloo", "jacket", "key", "lantern", "mirror", "necklace",
    "owl", "piano", "quiver", "rocket", "squirrel", "trophy", "unicorn", "vase", "window", "yacht",
    # ... Add more nouns to total 250

    # Adjectives (250)
    "brave", "calm", "delightful", "eager", "fancy", "gentle", "happy", "innocent", "jolly", "kind",
    "lively", "magnificent", "noble", "optimistic", "peaceful", "quick", "radiant", "shy", "tidy", "unique",
    "vivid", "warm", "yellow", "zealous", "adorable", "beautiful", "charming", "diligent", "energetic", "fierce",
    "graceful", "humble", "intelligent", "jovial", "keen", "lovely", "merry", "neat", "outstanding", "pleasant",
    "quirky", "respectful", "silly", "thoughtful", "upbeat", "vibrant", "whimsical", "youthful", "zany",
    # ... Add more adjectives to total 250

    # Verbs (250)
    "accept", "bounce", "climb", "dance", "explore", "fly", "gather", "help", "imagine", "jump",
    "kick", "laugh", "move", "notice", "open", "play", "question", "run", "sing", "talk",
    "understand", "visit", "wait", "yell", "zoom", "answer", "build", "create", "dig", "enjoy",
    "focus", "grow", "hunt", "identify", "juggle", "know", "learn", "measure", "negotiate", "observe",
    "perform", "quiet", "record", "search", "travel", "update", "volunteer", "wander", "write",
    # ... Add more verbs to total 250

    # Adverbs (250)
    "abruptly", "beautifully", "carefully", "diligently", "eagerly", "faithfully", "gracefully", "happily",
    "immediately", "joyfully", "kindly", "loudly", "magically", "neatly", "openly", "politely", "quickly",
    "rarely", "silently", "thoughtfully", "unexpectedly", "vividly", "warmly", "yawningly", "zealously",
    "accidentally", "boldly", "cheerfully", "deliberately", "enthusiastically", "frequently", "gently", "honestly",
    "intensely", "justly", "knowingly", "lightly", "merrily", "nervously", "officially", "partially", "quietly",
    "readily", "safely", "terribly", "urgently", "vaguely", "wildly", "yearly", "zestfully",
    # ... Add more adverbs to total 250
]


class WordGame:

    def __init__(self):
        self.points = 0
        self.target_word = ""
        self.attempts = 3
        self.generate_task()
    
    def generate_task(self):
        self.attempts = 3
        self.target_word = random.choice(words)
        print(f"New target word: {self.target_word}")

    def check_input_for_word(self, string):
        if self.target_word in string:
            print(f"The player input the target word and the task was reset.")
            self.generate_task()
        else:
            print(f"The player did not input the target word, so that's good.")
            

    def check_output_for_word(self, string):
        if self.target_word in string:
            self.points += self.attempts
            print(f"The player scored {self.attempts} points and has a total of {self.points}.")
            self.generate_task()
        else:
            print("The response did not contain the word.")
            self.attempts -= 1
            print(f"Remaining attempts: {self.attempts}")
            if self.attempts <= 0:
                self.generate_task()
                print(f"The player ran out of attempts.")
            else:
                print(f"The player has attempts left.")
    

game = WordGame()



def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    print("message:")
    print(message)
    game.check_input_for_word(message)

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response
    print("response:")
    print(response)
    game.check_output_for_word(response)


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond,
    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)",
        ),
    ],
)


if __name__ == "__main__":
    demo.launch()