import os
import json
import shutil
import gradio as gr
from huggingface_hub import snapshot_download

def process_and_zip_folders(huggingface_dataset_url, output_dir):
    # データセットをダウンロード
    dataset_path = snapshot_download(huggingface_dataset_url)

    # 処理対象のフォルダ
    folders = [
        "hololive-jp", "hololive-id", "hololive-en", "holostars",
        "hololive-cn", "nijisanji", "nijisanji-en", "vshojo",
        "phaseconnect", "indies", "other"
    ]

    zip_files = []

    for folder in folders:
        folder_path = os.path.join(dataset_path, folder)
        if not os.path.exists(folder_path):
            continue

        # JSONファイルを探す
        json_path = os.path.join(folder_path, "model_info.json")
        if not os.path.exists(json_path):
            continue

        # JSONを読み込む
        with open(json_path, "r", encoding="utf-8") as f:
            model_info = json.load(f)

        # ZIPに圧縮
        zip_name = f"{folder}.zip"
        zip_path = os.path.join(output_dir, zip_name)
        shutil.make_archive(base_name=zip_path.replace('.zip', ''), format="zip", root_dir=folder_path)
        zip_files.append(zip_path)

    return zip_files

def gradio_interface():
    def start_process(huggingface_url, output_directory):
        zip_files = process_and_zip_folders(huggingface_url, output_directory)
        return zip_files

    interface = gr.Interface(
        fn=start_process,
        inputs=[
            gr.Textbox(label="Hugging Face Dataset URL", placeholder="https://huggingface.co/datasets/soiz1/rvc-models"),
            gr.Textbox(label="Output Directory", placeholder="/path/to/output")
        ],
        outputs=gr.File(label="Generated ZIP Files"),
        title="Folder to ZIP Generator",
        description="指定されたフォルダを取得してZIPに圧縮します。"
    )

    interface.launch()

if __name__ == "__main__":
    gradio_interface()