import gradio as gr
import zipfile
import os
import shutil
from datasets import load_dataset

def zip_folder(dataset_name, folder_name, output_zip):
    """
    指定されたHugging Faceデータセット内のフォルダをZIPに圧縮する。

    Args:
        dataset_name (str): データセット名。
        folder_name (str): 圧縮するフォルダの名前。
        output_zip (str): 出力するZIPファイル名。

    Returns:
        str: ZIPファイルのパス。
    """
    # データセットをロード
    dataset = load_dataset(dataset_name)

    # 一時ディレクトリに保存
    temp_dir = "temp_dataset"
    if os.path.exists(temp_dir):
        shutil.rmtree(temp_dir)
    os.makedirs(temp_dir)

    # データを保存
    dataset.save_to_disk(temp_dir)

    # フォルダパスを指定
    folder_path = os.path.join(temp_dir, folder_name)
    if not os.path.exists(folder_path):
        return f"指定されたフォルダ {folder_name} が見つかりませんでした。"

    # ZIPファイルに圧縮
    with zipfile.ZipFile(output_zip, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(folder_path):
            for file in files:
                file_path = os.path.join(root, file)
                arcname = os.path.relpath(file_path, start=folder_path)
                zipf.write(file_path, arcname)

    # 一時ディレクトリを削除
    shutil.rmtree(temp_dir)

    return output_zip

def gradio_interface(dataset_name, folder_name):
    """
    Gradioインターフェース用関数。

    Args:
        dataset_name (str): データセット名。
        folder_name (str): フォルダ名。

    Returns:
        file: ZIPファイル。
    """
    output_zip = "output.zip"
    result = zip_folder(dataset_name, folder_name, output_zip)

    if not os.path.exists(result):
        return f"エラー: {result}"

    return result

# Gradio UI
description = "Hugging Faceデータセット内の指定フォルダをZIP形式でダウンロード"
iface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Textbox(label="Hugging Faceデータセット名", placeholder="例: squad"),
        gr.Textbox(label="フォルダ名", placeholder="例: train"),
    ],
    outputs=gr.File(label="ダウンロード: ZIPファイル"),
    title="Hugging FaceフォルダZIP変換",
    description=description
)

if __name__ == "__main__":
    iface.launch()