"""
eden_upload_weights.py
Uploads .pth model weights to their respective HF repos.
Skips files already present on HF. Shows size + progress per file.
"""
from huggingface_hub import HfApi, create_repo, upload_file, list_repo_files
import os, glob

TOKEN = os.environ.get("HF_TOKEN", "")
USER  = "Shanmuk4622"
BASE  = os.path.dirname(os.path.abspath(__file__))

api = HfApi(token=TOKEN)
print(f"Logged in as: {api.whoami()['name']}\n")

def parse_arch_ds(fn):
    fn = fn.lower().replace("\\", "/")
    ds, arch = "unknown", "unknown"
    if   "cifar100"  in fn: ds = "CIFAR-100"
    elif "cifar10"   in fn: ds = "CIFAR-10"
    elif "imagenet"  in fn: ds = "Custom-ImageNet300"
    if   "efficientnet" in fn: arch = "EfficientNetV2"
    elif "convnext"     in fn: arch = "ConvNeXtV2"
    elif "mobilevit"    in fn: arch = "MobileViTv3"
    elif "resnet50"     in fn: arch = "ResNet50"
    elif "resnet18"     in fn: arch = "ResNet18"
    elif "vgg16"        in fn: arch = "VGG16"
    elif "alexnet"      in fn: arch = "AlexNet"
    elif "inception"    in fn: arch = "InceptionV3"
    elif "densenet"     in fn: arch = "DenseNet121"
    elif "unet"         in fn: arch = "UNet"
    return arch, ds

def mb(path):
    return os.path.getsize(path) / 1_048_576

def already_uploaded(repo_id, filename):
    try:
        files = list(list_repo_files(repo_id, token=TOKEN, repo_type="model"))
        return filename in files
    except:
        return False

# ── Collect all .pth files ────────────────────────────────────────────────────
pth_files = sorted(glob.glob(os.path.join(BASE, "**", "*.pth"), recursive=True))

# Build upload plan: (local_path, repo_id, filename)
plan = []
skipped_unknown = []

for pth in pth_files:
    rel  = os.path.relpath(pth, BASE)
    arch, ds = parse_arch_ds(rel)
    if arch == "unknown" or ds == "unknown":
        skipped_unknown.append(rel)
        continue
    repo_id  = f"{USER}/EDEN-{arch}-{ds.replace(' ', '-')}"
    filename = os.path.basename(pth)
    plan.append((pth, repo_id, filename, rel))

total_mb = sum(mb(p[0]) for p in plan)
print(f"Files to upload : {len(plan)}")
print(f"Total size      : {total_mb:.1f} MB")
print(f"Skipped (unknown arch/dataset): {len(skipped_unknown)}")
if skipped_unknown:
    for s in skipped_unknown:
        print(f"  - {s}")
print()

# ── Upload ────────────────────────────────────────────────────────────────────
ok, skip, fail = 0, 0, 0

for i, (local_path, repo_id, filename, rel) in enumerate(plan, 1):
    size = mb(local_path)
    prefix = f"[{i:02d}/{len(plan)}]  {filename}  ({size:.1f} MB)"

    # Check if already on HF
    if already_uploaded(repo_id, filename):
        print(f"  ⏭   {prefix}  →  already on HF, skipping")
        skip += 1
        continue

    try:
        print(f"  ⬆   {prefix}  →  {repo_id} ...")
        upload_file(
            path_or_fileobj=local_path,
            path_in_repo=filename,
            repo_id=repo_id,
            token=TOKEN,
            repo_type="model",
        )
        print(f"  ✓   DONE")
        ok += 1
    except Exception as e:
        print(f"  ✗   FAILED: {e}")
        fail += 1

print()
print("="*60)
print(f"WEIGHTS UPLOAD SUMMARY")
print(f"  Uploaded : {ok}")
print(f"  Skipped  : {skip}  (already on HF)")
print(f"  Failed   : {fail}")
print(f"  Check    : https://huggingface.co/{USER}")
print("="*60)