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

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: set HF_TOKEN=your_token")
        sys.exit(1)

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

    local_dir = os.path.dirname(os.path.abspath(__file__))

    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}"
    print(f"[HF] Uploading modified sliding-window strict-small scripts to '{repo_id}' main branch...")
    
    ignore_patterns = [
        "**/__pycache__/*",
        "**/*.pyc",
        "**/hf_cache/*",
        "**/nltk_data/*",
        "**/checkpoints/*"
    ]
    
    sensitive_ignores = []
    token_pattern = re.compile(r"hf_[a-zA-Z0-9]{34}")
    for root, dirs, files in os.walk(local_dir):
        if "__pycache__" in root or "checkpoints" in root:
            continue
        for file in files:
            file_path = os.path.join(root, file)
            if file.endswith(('.bin', '.safetensors', '.zip', '.tar.gz', '.pkl')):
                continue
            try:
                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: {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="Update strict-small architecture files (sliding window [64, 16, 8, 4], ln3, ln_post_moe, no res3)"
        )
        print(f"\n[HF] Success! Scripts 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)