# Standard library imports
import os
import sys
import argparse
import subprocess

# Third-party imports
import yaml

# Parse command-line arguments
parser = argparse.ArgumentParser(description="Download files from Hugging Face Hub with authentication.")
parser.add_argument("repo_id", help="Repository ID (e.g. username/repo-name)")
parser.add_argument("--quant", type=str, help="Quantization extension to include (e.g. IQ2_XS, q4_0, etc.)")
args = parser.parse_args()

repo_id = args.repo_id
quant = args.quant

if quant:
    print(f"Quant mode: {quant}")
else:
    print("Quant mode: disabled")

with open('config.yaml', 'r') as f:
    config = yaml.safe_load(f)

default_location = config['default_location']
token = config.get('token')  # Optional: store token in config.yaml

safe_repo_name = repo_id.replace('/', '_').replace('\\', '_')
download_location = os.path.join(default_location, safe_repo_name)
os.makedirs(download_location, exist_ok=True)

default_location = config['default_location']
token = config.get('token')  # Optional: store token in config.yaml

# Authenticate with Hugging Face CLI
if token:
    try:
        subprocess.run(['hf', 'auth', 'login', '--token', token], check=True)
    except subprocess.CalledProcessError as e:
        print("\n[ERROR] Authentication failed. Please check your Hugging Face token in config.yaml.\n"
              "If you want to add the token to the git credential helper, set 'add_to_git_credential' to True in your config or use the '--add-to-git-credential' flag.\n"
              "Original error: Invalid user token or authentication failed.\n")
        exit(1)
else:
    print("No token found in config.yaml. Please add your Hugging Face token.")
    exit(1)




# Download README.md and config.json (no --include, just as positional arguments)
download_cmd_1 = [
    'hf', 'download', '--local-dir', download_location, repo_id, 'README.md', 'config.json'
]
print("Download command 1:")
print(' '.join(download_cmd_1))
subprocess.run(download_cmd_1, check=True)
print(f"Downloaded ['README.md', 'config.json'] to {download_location}")


# Second hf download call: use --include for quantized files
if quant:
    quant_pattern = f"*{quant}.*"
    download_cmd_2 = [
        'hf', 'download', '--local-dir', download_location, repo_id,
        '--include', quant_pattern
    ]
    print("Download command 2:")
    print(' '.join(download_cmd_2))
    subprocess.run(download_cmd_2, check=True)