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

def zip_folder(dataset_name, folder_name, output_zip):
    try:
        print(f"ロード対象データセット: {dataset_name}")
        dataset = load_dataset(dataset_name)
    except Exception as e:
        return f"データセットをロードできませんでした: {str(e)}"

    if not hasattr(dataset, "save_to_disk"):
        return f"データセット '{dataset_name}' は save_to_disk に対応していません。"

    temp_dir = f"temp_dataset_{uuid.uuid4().hex}"
    os.makedirs(temp_dir, exist_ok=True)

    try:
        dataset.save_to_disk(temp_dir)
        folder_path = os.path.join(temp_dir, folder_name)

        if not os.path.exists(folder_path):
            available_folders = os.listdir(temp_dir)
            return f"指定されたフォルダ '{folder_name}' は見つかりません。利用可能なフォルダ: {available_folders}"

        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)

    except Exception as e:
        return f"エラーが発生しました: {str(e)}"
    finally:
        shutil.rmtree(temp_dir)

    return output_zip

def gradio_interface(dataset_name, folder_name):
    output_zip = "output.zip"
    result = zip_folder(dataset_name, folder_name, output_zip)

    if not os.path.exists(output_zip):
        return result

    return output_zip

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.Textbox(label="結果"),  # エラーメッセージ用
        gr.File(label="ダウンロード: ZIPファイル"),
    ],
    title="Hugging FaceフォルダZIP変換",
    description=description
)

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