import os
import zipfile
import datasets
import gradio as gr

# Hugging Faceのデータセットをダウンロードして指定されたフォルダをZIP圧縮する関数
def download_and_zip(repo_id, folder_name):
    try:
        # データセットをHugging Face Hubからダウンロード
        dataset = datasets.load_dataset(repo_id)

        # ダウンロードされたデータセットのパスを取得
        dataset_path = dataset.data_files['train']  # 'train' データセットを例にします。必要に応じて変更

        # フォルダのパス
        folder_path = os.path.join(dataset_path, folder_name)

        # ZIP圧縮の保存先パス
        zip_path = f"{folder_name}.zip"

        # フォルダが存在するかをチェック
        if os.path.exists(folder_path) and os.path.isdir(folder_path):
            # フォルダをZIP圧縮する
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for root, dirs, files in os.walk(folder_path):
                    for file in files:
                        zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), folder_path))
            return zip_path
        else:
            return "指定されたフォルダが見つかりません"
    except Exception as e:
        return f"エラー: {str(e)}"

# Gradioインターフェースの作成
def interface(repo_id, folder_name):
    return download_and_zip(repo_id, folder_name)

# Gradioインターフェースのセットアップ
gr.Interface(
    fn=interface,
    inputs=["text", "text"],
    outputs="file",
    live=True,
    description="Hugging Faceのデータセットの指定フォルダをZIPに圧縮します。"
).launch()