# app.py
import gradio as gr
from datasets import load_dataset

# Загружаем только 500 случайных изображений (можно увеличить до 1000–2000)
dataset = load_dataset("deepghs/sankaku_full", split="train[:500]")

def search_images(tag):
    results = []
    for item in dataset:
        tags = item.get("tags", "")
        if tag.lower() in tags.lower():
            # Добавляем миниатюру + теги
            results.append((item["image"], f"Tags: {tags}"))
            if len(results) >= 10:  # показываем максимум 10 результатов
                break
    return results if results else [("https://via.placeholder.com/300?text=No+Results", "Ничего не найдено")]

with gr.Blocks(title="Sankaku Viewer") as demo:
    gr.Markdown("# 🔞 Sankaku Full Dataset Viewer")
    gr.Markdown("⚠️ **NSFW content. For adults only.**")
    
    tag_input = gr.Textbox(label="Поиск по тегу (например: wolf, furry, animal ears)")
    gallery = gr.Gallery(label="Результаты", columns=3, rows=4, height=600)

    tag_input.change(fn=search_images, inputs=tag_input, outputs=gallery)

demo.launch()