import os
from huggingface_hub import hf_hub_download, list_repo_files
import tarfile
import shutil
from tqdm import tqdm

def download_and_extract_file(repo_id, filename, local_dir):
    """
    Download and extract a file from Hugging Face.
    
    Args:
        repo_id (str): The repository ID on Hugging Face
        filename (str): The name of the file to download
        local_dir (str): The directory to extract files to
    """
    print(f"Starting download of {filename}...")
    
    # Create directory if it doesn't exist
    os.makedirs(local_dir, exist_ok=True)
    
    try:
        # Download the file using hf_hub_download
        file_path = hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            repo_type="space",
            local_dir=local_dir
        )
        
        print(f"Download completed. File saved to: {file_path}")
        
        # Extract the tar.xz file with progress bar
        print(f"Extracting {filename}...")
        with tarfile.open(file_path, 'r:xz') as tar:
            # Get total number of files
            total_files = len(tar.getmembers())
            
            # Extract with progress bar
            for member in tqdm(tar.getmembers(), total=total_files, desc=f"Extracting {filename}"):
                tar.extract(member, path=local_dir)
        
        print(f"Extraction of {filename} completed successfully!")
        
        # Clean up the downloaded archive
        os.remove(file_path)
        print(f"Cleaned up downloaded archive {filename}.")
        
    except Exception as e:
        print(f"An error occurred while processing {filename}: {str(e)}")
        # Clean up in case of error
        if os.path.exists(file_path):
            os.remove(file_path)
        raise

def download_video_files(repo_id, local_dir):
    """
    Download all video files from the downloads directory.
    
    Args:
        repo_id (str): The repository ID on Hugging Face
        local_dir (str): The directory to save videos to
    """
    print("Starting download of video files...")
    
    # Create downloads directory if it doesn't exist
    os.makedirs(local_dir, exist_ok=True)
    
    try:
        # List all files in the downloads directory
        files = list_repo_files(
            repo_id=repo_id,
            repo_type="space",
            revision="main"
        )
        
        # Filter for video files in the downloads directory
        video_files = [f for f in files if f.startswith('downloads/') and f.endswith('.mp4')]
        
        if not video_files:
            print("No video files found in the downloads directory.")
            return
        
        print(f"Found {len(video_files)} video files to download.")
        
        # Download each video file
        for video_file in video_files:
            print(f"\nDownloading {video_file}...")
            hf_hub_download(
                repo_id=repo_id,
                filename=video_file,
                repo_type="space",
                local_dir=local_dir
            )
            print(f"Downloaded {video_file} successfully!")
        
        print("\nAll video files downloaded successfully!")
        
    except Exception as e:
        print(f"An error occurred while downloading video files: {str(e)}")
        raise

def download_and_extract_precomputed():
    """
    Download and extract precomputed results from Hugging Face.
    Downloads both intermediate_results.tar.xz and outputs.tar.xz,
    and downloads all video files from the downloads directory.
    """
    repo_id = "abdullahyang/NonisotropicSkeletonDiffusion_PrecomputedResults"
    
    # Download and extract intermediate_results.tar.xz
    download_and_extract_file(
        repo_id=repo_id,
        filename="intermediate_results.tar.xz",
        local_dir="."
    )
    
    # Download and extract outputs.tar.xz
    download_and_extract_file(
        repo_id=repo_id,
        filename="outputs.tar.xz",
        local_dir="."
    )
    
    # Download video files
    download_video_files(
        repo_id=repo_id,
        local_dir="."
    )
    
    print("All downloads and extractions completed successfully!")

if __name__ == "__main__":
    download_and_extract_precomputed()