"""
Wan2.1-VACE-1.3Bモデルのダウンロードユーティリティ
"""
import os
import logging
from pathlib import Path
from huggingface_hub import snapshot_download, hf_hub_download
from tqdm import tqdm

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

MODEL_ID = "Wan-AI/Wan2.1-VACE-1.3B"
DEFAULT_MODEL_DIR = "./Wan2.1-VACE-1.3B"

def download_model(model_dir=DEFAULT_MODEL_DIR, force=False):
    """
    Wan2.1-VACE-1.3Bモデルをダウンロード
    
    Args:
        model_dir: モデルを保存するディレクトリ
        force: 既存のモデルを強制的に再ダウンロード
    
    Returns:
        str: ダウンロードしたモデルのディレクトリパス
    """
    model_path = Path(model_dir)
    
    # 既にモデルが存在する場合
    if model_path.exists() and not force:
        logger.info(f"モデルは既に存在します: {model_path}")
        return str(model_path)
    
    logger.info(f"モデルをダウンロード中: {MODEL_ID}")
    logger.info(f"保存先: {model_path}")
    
    try:
        # モデル全体をダウンロード
        snapshot_download(
            repo_id=MODEL_ID,
            local_dir=model_dir,
            resume_download=True,
            tqdm_class=tqdm
        )
        logger.info("モデルのダウンロードが完了しました")
        return str(model_path)
        
    except Exception as e:
        logger.error(f"モデルのダウンロード中にエラーが発生しました: {e}")
        raise

def check_model_files(model_dir=DEFAULT_MODEL_DIR):
    """
    必要なモデルファイルが存在するかチェック
    
    Returns:
        bool: すべての必要ファイルが存在する場合True
    """
    model_path = Path(model_dir)
    
    # 必要なファイルのリスト（実際のファイル名は要確認）
    required_files = [
        "diffusion_pytorch_model.safetensors",
        "Wan2.1_VAE.pth",
        "models_t5_umt5-xxl-enc-bf16.pth",
        "config.json"
    ]
    
    missing_files = []
    for file in required_files:
        if not (model_path / file).exists():
            missing_files.append(file)
    
    if missing_files:
        logger.warning(f"不足しているファイル: {missing_files}")
        return False
    
    return True

if __name__ == "__main__":
    # スクリプトを直接実行した場合、モデルをダウンロード
    download_model()