import os
import gradio as gr
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage

# Function to handle responses
def respond(user_input, history):
    # Configure Hugging Face model endpoint
    llm = HuggingFaceEndpoint(
        repo_id="HuggingFaceTB/SmolLM2-1.7B-Instruct",
        task="text-generation",
        temperature=0.5,
        repetition_penalty=1.03,
        huggingfacehub_api_token=os.getenv('HUGGING_FACE_TOKEN')
    )
    chat = ChatHuggingFace(llm=llm, verbose=True)

    # Prepare system message to set the context
    system_message = "You are playing a word guessing game. Guess ONE word based on previous guesses and avoid repeats. DO NOT ANSWER WITH MORE THAN ONE WORD."

    # Track guessed words from history
    last_human_message, last_ai_message = history[-1] if history else ("", "")

    # Add messages to the conversation
    messages = [
        SystemMessage(content=system_message),
    ]

    if last_human_message:
        messages.append(SystemMessage(content="Pick a word that is similar to both " + last_human_message + " and " + last_ai_message))
    else:
        messages.append(SystemMessage(content="First message. Pick a random word!"))

    # Load past conversation into messages
    # for human, ai in history:
    #     messages.append(HumanMessage(content=human))
    #     messages.append(AIMessage(content=ai))

    # Initialize response as an empty string
    response = ""

    # Stream response from the model
    for message in chat.stream(messages):
        if isinstance(message, str):
            response += message
        else:
            response += message.content

        yield response

# Set up the Gradio interface
demo = gr.ChatInterface(
    respond,
    title="WordSync",
    description="A word guessing game where you and the AI try to guess the same word. The AI remembers past guesses!"
)

demo.launch(share=True)