from datasets import load_dataset
import gradio as gr
import random
import requests
from PIL import Image
from io import BytesIO

# Load the dataset
steam_dataset = load_dataset("taesiri/SteamCommunityImagesDailyDumpy")

# Convert to pandas once and keep it in memory
df = steam_dataset["train"].to_pandas()
game_counts = df["game_name"].value_counts().to_dict()
game_list = list(game_counts.keys())

# Cache for storing processed game data
game_cache = {}


def get_random_image(game_name):
    # Lazy loading: only process game data when first requested
    if game_name not in game_cache:
        game_cache[game_name] = df[df["game_name"] == game_name]["image_url"].tolist()

    # Get random image URL from cached list
    random_url = random.choice(game_cache[game_name])

    # Download and return image
    response = requests.get(random_url)
    img = Image.open(BytesIO(response.content))
    return img, f"Total images available for {game_name}: {game_counts[game_name]}"


# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Steam Game Image Preview")

    with gr.Row():
        game_dropdown = gr.Dropdown(
            choices=game_list,
            label="Select a Game",
            info="Choose a game to see random screenshots",
        )

    with gr.Row():
        show_btn = gr.Button("Show Random Image")

    with gr.Column():
        image_output = gr.Image(type="pil", label="Random Game Image")
        stats_output = gr.Textbox(label="Statistics")

    show_btn.click(
        fn=get_random_image,
        inputs=[game_dropdown],
        outputs=[image_output, stats_output],
    )

demo.launch()