import gradio as gr
import os
from huggingface_hub import InferenceClient

api_token = os.getenv('token')

client = InferenceClient("gpt2", token=api_token)

scenario = "You wake up in a strange forest. Paths lead in all directions."
choices = [
    "Walk north into the dense forest.",
    "Go south back towards a distant town.",
    "Explore east towards the sound of water.",
    "Move west where the land rises."
]

history = []

def generate_continuation(choice):
    global scenario, history
    history.append(choice)
    prompt = f"{scenario} {' '.join(history)}"
    max_length = 50  # Adjust as needed

    response = client(prompt, max_length=max_length, remove_input=True)
    continuation = response[0]['generated_text']

    return continuation

with gr.Blocks() as app:
    with gr.Row():
        story_display = gr.Textbox(value=scenario, lines=7)
        choice_buttons = [gr.Button(choice) for choice in choices]
    
    def handle_choice(choice_index):
        chosen_action = choices[choice_index]
        new_scenario = generate_continuation(chosen_action)
        story_display.update(value=new_scenario)
        return gr.update()

    for i, button in enumerate(choice_buttons):
        button.click(fn=lambda x=i: handle_choice(x), inputs=None, outputs=story_display)

app.launch()