import streamlit as st
from huggingface_hub import InferenceClient
import random

# Replace with your desired model ID from Hugging Face Hub
client = InferenceClient("unsloth/Llama-3.2-1B-Instruct") 

words = [
    "apple", "pear", "banana", "strawberry", "mad"
]


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)
        st.session_state["target_word"] = self.target_word  # Store in session state
        st.title(f'The Game - Current score: {self.points}, remaining attempts: {self.attempts}')

    def check_input_for_word(self, user_input):
        if self.target_word in user_input.lower():
            self.generate_task()
            self.points -= 1
            st.write(f"You input the target word yourself, so you lost one point and the game reset.")
            return True
        else:
            return False

    def check_output_for_word(self, response):
        if self.target_word in response.lower():
            self.points += self.attempts
            score_gained = self.attempts
            self.generate_task()
            return f"Success! You earned {score_gained} points!"
        else:
            self.attempts -= 1
            if self.attempts <= 0:
                self.generate_task()
                return f"You did not win in three attempts. Generating new target word."
            else:
                return "That didn't quite hit the mark. Try again!"

    def respond(self, user_message):
        messages = [{"role": "system", "content": st.session_state["system_message"]}]
        
        for val in st.session_state["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": user_message})
        response = ""
        for message in client.chat_completion(
            messages, max_tokens=st.session_state["max_tokens"], stream=True,
            temperature=st.session_state["temperature"], top_p=st.session_state["top_p"]
        ):
            token = message.choices[0].delta.content
            response += token
        
        output_check_result = self.check_output_for_word(response)
        st.write(response + f"\n\n---\n\n{output_check_result}")
        st.write(f'Current score: {self.points}, remaining attempts: {self.attempts}')

        # Update session state with history and system message
        st.session_state["history"].append((user_message, response))
        st.session_state["system_message"] = response


game = WordGame()

# Session state variables to store dynamic information
st.session_state["history"] = []  # List of tuples (user message, response)
st.session_state["system_message"] = "You are a friendly Chatbot."  # Initial system message


# Update session state with user input for system message
st.session_state["system_message"] = system_message_input

user_message = st.text_