import gradio as gr
import requests
import subprocess
import zipfile
import os
import shutil
from huggingface_hub import snapshot_download, login, hf_hub_download

def clean_workspace():
    for item in os.listdir():
        if os.path.isdir(item) or item.endswith(".zip"):
            try:
                if os.path.isdir(item):
                    shutil.rmtree(item)
                else:
                    os.remove(item)
            except Exception as e:
                print(f"削除失敗: {item} ({e})")

def zip_path_from_folder(folder_path, zip_path):
    """指定フォルダをZIP化"""
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        if os.path.isdir(folder_path):
            for root, dirs, files in os.walk(folder_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    zipf.write(file_path, os.path.relpath(file_path, folder_path))
        elif os.path.isfile(folder_path):
            zipf.write(folder_path, os.path.basename(folder_path))
        else:
            raise ValueError(f"指定フォルダが存在しません: {folder_path}")
    return zip_path

def download_and_zip_git_clone(repo_url, target_folder=None):
    repo_name = repo_url.rstrip('/').split('/')[-1]
    subprocess.run(["git", "clone", repo_url], check=True)
    
    zip_name = f"{repo_name}.zip" if not target_folder else f"{repo_name}_{target_folder.replace('/', '_')}.zip"
    folder_to_zip = os.path.join(repo_name, target_folder) if target_folder else repo_name
    
    zip_path = zip_path_from_folder(folder_to_zip, zip_name)
    
    shutil.rmtree(repo_name)
    return zip_path

def download_and_zip_github_api(repo_url, target_folder=None):
    repo_name = repo_url.rstrip('/').split('/')[-1]
    zip_url = f"{repo_url}/archive/refs/heads/main.zip"
    response = requests.get(zip_url)
    if response.status_code != 200:
        raise Exception(f"GitHub API ダウンロード失敗: ステータスコード {response.status_code}")

    temp_zip = f"{repo_name}_full.zip"
    with open(temp_zip, 'wb') as f:
        f.write(response.content)
    
    # ZIP展開して特定フォルダを再ZIP化
    temp_dir = f"{repo_name}_temp"
    with zipfile.ZipFile(temp_zip, 'r') as zip_ref:
        zip_ref.extractall(temp_dir)
    os.remove(temp_zip)
    
    folder_to_zip = os.path.join(temp_dir, f"{repo_name}-main", target_folder) if target_folder else os.path.join(temp_dir, f"{repo_name}-main")
    zip_name = f"{repo_name}.zip" if not target_folder else f"{repo_name}_{target_folder.replace('/', '_')}.zip"
    
    zip_path = zip_path_from_folder(folder_to_zip, zip_name)
    shutil.rmtree(temp_dir)
    return zip_path

def download_and_zip_huggingface_hub(repo_url, token=None, specific_file=None, target_folder=None):
    repo_id = '/'.join(repo_url.rstrip('/').split('/')[-2:])
    if token:
        login(token=token)
    
    if specific_file:
        specific_path = hf_hub_download(repo_id=repo_id, filename=specific_file)
        zip_name = f"{repo_id.replace('/', '_')}_{specific_file.replace('/', '_')}.zip"
        zip_path = zip_path_from_folder(specific_path, zip_name)
    else:
        # 全体をsnapshot_downloadで取得してフォルダ指定でZIP化
        folder_path = snapshot_download(repo_id=repo_id, use_auth_token=token)
        folder_to_zip = os.path.join(folder_path, target_folder) if target_folder else folder_path
        zip_name = f"{repo_id.replace('/', '_')}.zip" if not target_folder else f"{repo_id.replace('/', '_')}_{target_folder.replace('/', '_')}.zip"
        zip_path = zip_path_from_folder(folder_to_zip, zip_name)
    return zip_path

def process_repo(repo_url, method, use_token, token, specific_file, download_type, target_folder):
    try:
        clean_workspace()
        token_value = token if use_token else None

        if method == "git clone":
            zip_path = download_and_zip_git_clone(repo_url, target_folder)
        elif method == "GitHub API":
            zip_path = download_and_zip_github_api(repo_url, target_folder)
        elif method == "Hugging Face Hub":
            if download_type == "特定のファイルのみ":
                zip_path = download_and_zip_huggingface_hub(repo_url, token_value, specific_file)
            else:
                zip_path = download_and_zip_huggingface_hub(repo_url, token_value, None, target_folder)
        else:
            return f"不明な方法: {method}", None, gr.update(visible=False)

        return f"成功: {zip_path}", zip_path, gr.update(visible=True)
    except Exception as e:
        return f"エラー: {str(e)}", None, gr.update(visible=False)

# Gradio UI
with gr.Blocks() as app:
    gr.Markdown("# リポジトリZIP変換ツール")
    repo_url_input = gr.Textbox(label="リポジトリURL")
    method_input = gr.Radio(["git clone", "GitHub API", "Hugging Face Hub"], label="ダウンロード方法")
    use_token_input = gr.Checkbox(label="トークンを使用")
    token_input = gr.Textbox(label="Hugging Faceトークン", visible=False)
    download_type_input = gr.Radio(["リポジトリ全体", "特定のファイルのみ"], label="ダウンロードタイプ", visible=False)
    specific_file_input = gr.Textbox(label="特定ファイルのパス", visible=False)
    target_folder_input = gr.Textbox(label="特定フォルダのみZIP化（任意）")
    output_text = gr.Textbox(label="結果")
    download_button = gr.Button("ZIPに変換してダウンロード")
    download_file = gr.File(label="ZIPファイル", visible=False)

    # 表示切替
    use_token_input.change(lambda x: gr.update(visible=x), inputs=use_token_input, outputs=token_input)
    method_input.change(lambda x: gr.update(visible=(x=="Hugging Face Hub")), inputs=method_input, outputs=download_type_input)
    download_type_input.change(lambda x: gr.update(visible=(x=="特定のファイルのみ")), inputs=download_type_input, outputs=specific_file_input)

    download_button.click(
        fn=process_repo,
        inputs=[repo_url_input, method_input, use_token_input, token_input, specific_file_input, download_type_input, target_folder_input],
        outputs=[output_text, download_file, download_file],
    )

app.launch()