#!/usr/bin/env python3
import os
import time
from huggingface_hub import HfApi, list_repo_files

# ==============================================================
# CONFIGURATION
# ==============================================================

REPO_ID = "chichimedia/Greywan"
HF_TOKEN = os.getenv("HF_TOKEN")

if not HF_TOKEN:
    raise EnvironmentError("HF_TOKEN environment variable not set!")

api = HfApi()

# ==============================================================
# HELPERS
# ==============================================================

def list_local_files(base_dir="."):
    """Recursively list all local files excluding system/git/junk."""
    excluded_dirs = {
        ".git", "__pycache__", ".cache", ".idea", ".vscode", ".venv", "env",
        "venv", "node_modules", ".mypy_cache"
    }
    excluded_exts = {".lock", ".tmp", ".log", ".bat", ".sh"}

    all_files = []
    for root, dirs, files in os.walk(base_dir):
        dirs[:] = [d for d in dirs if d not in excluded_dirs]
        for f in files:
            if any(f.endswith(ext) for ext in excluded_exts):
                continue
            full_path = os.path.join(root, f)
            rel_path = os.path.relpath(full_path, base_dir)
            all_files.append(rel_path.replace("\\", "/"))
    return all_files

# ==============================================================
# MAIN UPLOAD LOGIC
# ==============================================================

def main():
    print("[INFO] Fetching existing files from HF repo...")
    remote_files = set(list_repo_files(REPO_ID, token=HF_TOKEN))
    print(f"[INFO] Found {len(remote_files)} files on remote.\n")

    local_files = list_local_files(".")
    upload_count = 0
    skip_count = 0

    for rel in sorted(local_files):
        if rel in remote_files:
            print(f"[SKIP] {rel}")
            skip_count += 1
            continue

        print(f"[UPLOAD] {rel}")
        full = os.path.abspath(rel)
        start_time = time.time()

        try:
            api.upload_file(
                path_or_fileobj=full,
                path_in_repo=rel,
                repo_id=REPO_ID,
                repo_type="model",
                token=HF_TOKEN,
            )
            elapsed = time.time() - start_time
            size_mb = os.path.getsize(full) / (1024 * 1024)
            speed = size_mb / elapsed if elapsed > 0 else 0
            print(f"[OK] {rel} — {size_mb:.1f} MB in {elapsed:.1f}s ({speed:.2f} MB/s real)\n")
            upload_count += 1
        except Exception as e:
            print(f"[ERROR] Failed to upload {rel}: {e}\n")

    print(f"\n[SUMMARY] Uploaded: {upload_count}, Skipped: {skip_count}")
    print("[DONE]")

# ==============================================================
# ENTRY POINT
# ==============================================================

if __name__ == "__main__":
    main()