import os
import sys
import re
import argparse
from huggingface_hub import HfApi, create_repo

def upload_project(repo_name=None):
    token = os.environ.get("HF_TOKEN")
    if not token:
        print("Error: HF_TOKEN environment variable not set!")
        print("Please set it before running: export HF_TOKEN=your_token")
        sys.exit(1)

    if not repo_name:
        repo_name = os.environ.get("HF_REPO_NAME", "Normal_c1")

    # Set local_dir to the directory containing this script
    local_dir = os.path.dirname(os.path.abspath(__file__))
    print(f"[HF] Source directory: {local_dir}")

    api = HfApi(token=token)
    try:
        user_info = api.whoami()
        username = user_info["name"]
        print(f"[HF] Authenticated as: {username}")
    except Exception as e:
        print(f"[HF] Authentication failed: {e}")
        sys.exit(1)

    repo_id = f"{username}/{repo_name}"
    
    # Auto-create the repo if it does not exist
    try:
        create_repo(repo_id=repo_id, repo_type="model", token=token, exist_ok=True)
        print(f"[HF] Repository '{repo_id}' is ready.")
    except Exception as e:
        print(f"[HF] Warning during repository creation: {e}")

    print(f"[HF] Uploading ALL project files (including checkpoints, results, scripts) to '{repo_id}' main branch...")
    
    # Define ignore patterns (only excluding huge library download caches)
    ignore_patterns = [
        "**/__pycache__/**",
        "**/*.pyc",
        "**/hf_cache/**",
        "**/nltk_data/**",
        "hf_cache/**",
        "nltk_data/**",
        "hf_cache",
        "nltk_data",
        "**/hf_cache/*",
        "**/nltk_data/*"
    ]
    
    sensitive_ignores = []
    token_pattern = re.compile(r"hf_[a-zA-Z0-9]{34}")
    for root, dirs, files in os.walk(local_dir):
        # Skip scanning cache directories for sensitive tokens
        if any(d in root for d in ["__pycache__", "hf_cache", "nltk_data"]):
            continue
        for file in files:
            file_path = os.path.join(root, file)
            # Skip checking binary or archive files for tokens to keep scan fast
            if file.endswith(('.bin', '.safetensors', '.zip', '.tar.gz', '.pkl', '.pt', '.pth')):
                continue
            try:
                # Skip checking files larger than 10MB
                if os.path.getsize(file_path) > 10 * 1024 * 1024:
                    continue
                with open(file_path, "r", errors="ignore") as f:
                    content = f.read()
                    if token_pattern.search(content):
                        rel_path = os.path.relpath(file_path, local_dir)
                        rel_path_glob = rel_path.replace("\\", "/")
                        sensitive_ignores.append(rel_path_glob)
                        print(f"  -> Warning: Sensitive file containing raw token excluded from upload: {rel_path_glob}")
            except Exception:
                pass

    all_ignores = ignore_patterns + sensitive_ignores

    try:
        api.upload_folder(
            folder_path=local_dir,
            repo_id=repo_id,
            repo_type="model",
            revision="main",
            ignore_patterns=all_ignores,
            commit_message="Upload full run contents (scripts, checkpoints, and results)"
        )
        print(f"\n[HF] Success! All files uploaded to: https://huggingface.co/{repo_id}/tree/main")
    except Exception as e:
        print(f"[HF] Upload failed: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo-name", type=str, default=None, help="Hugging Face repository name")
    args = parser.parse_args()
    upload_project(args.repo_name)