"""Convert OpenWorldSAM checkpoint to safetensors and upload to HuggingFace.

Run after `huggingface-cli login`:
    python convert_and_upload.py

The script is idempotent: if model.safetensors already exists, it skips conversion.
"""

import json
import os
import sys

import torch
from safetensors.torch import save_file

HF_REPO_DIR = os.path.dirname(os.path.abspath(__file__))
HF_REPO_ID = "neerajaabhyankar/openworld-sam"
CHECKPOINT_PATH = os.path.join(HF_REPO_DIR, "weights", "openworld_sam_ade20k.pt")
SAFETENSORS_PATH = os.path.join(HF_REPO_DIR, "model.safetensors")
CONFIG_PATH = os.path.join(HF_REPO_DIR, "config.json")


# Keys to drop before saving:
#   visual_model.* / mm_extractor.*           — top-level duplicates of evf_sam2.visual_model/mm_extractor
#   evf_sam2.text_hidden_fcs.*               — duplicate of top-level text_hidden_fcs.* (shared memory)
SKIP_PREFIXES = ("visual_model.", "mm_extractor.", "evf_sam2.text_hidden_fcs.")


def convert():
    print(f"Loading checkpoint from {CHECKPOINT_PATH} ...")
    if not os.path.exists(CHECKPOINT_PATH):
        print(f"ERROR: checkpoint not found at {CHECKPOINT_PATH}", file=sys.stderr)
        sys.exit(1)

    ckpt = torch.load(CHECKPOINT_PATH, map_location="cpu", weights_only=False)
    state_dict = ckpt.get("model", ckpt)

    before = len(state_dict)
    filtered = {
        k: v.contiguous() for k, v in state_dict.items()
        if not any(k.startswith(p) for p in SKIP_PREFIXES)
    }
    after = len(filtered)
    print(f"  {before} keys → {after} keys after stripping top-level duplicates")

    # Clone every tensor to break shared-memory aliases before serialising
    clean = {k: v.clone() for k, v in filtered.items()}

    print(f"Saving model.safetensors ...")
    save_file(clean, SAFETENSORS_PATH)
    print(f"  Saved → {SAFETENSORS_PATH}")


def write_config():
    # Import after setting up sys.path so relative imports work
    sys.path.insert(0, HF_REPO_DIR)
    from configuration_openworld_sam import OpenWorldSAMConfig

    config = OpenWorldSAMConfig()
    cfg_dict = config.to_dict()
    cfg_dict["auto_map"] = {
        "AutoConfig": "configuration_openworld_sam.OpenWorldSAMConfig",
        "AutoModel": "modeling_openworld_sam.OpenWorldSAMModel",
    }
    with open(CONFIG_PATH, "w") as f:
        json.dump(cfg_dict, f, indent=2)
    print(f"Wrote config.json")


def upload():
    from huggingface_hub import HfApi

    api = HfApi()
    print(f"Creating / verifying repo {HF_REPO_ID} ...")
    api.create_repo(HF_REPO_ID, repo_type="model", private=False, exist_ok=True)

    print(f"Uploading {HF_REPO_DIR} → {HF_REPO_ID} ...")
    api.upload_folder(
        folder_path=HF_REPO_DIR,
        repo_id=HF_REPO_ID,
        repo_type="model",
        ignore_patterns=[
            "weights/*",
            "*.pt",
            "*.pth",
            "**/__pycache__/*",
            "**/*.pyc",
            ".git/*",
            "**/*.so",
            "build/*",
            "**/*.egg-info/*",
            ".DS_Store",
            # scratch / experiment files
            "expt*.py",
            "test_*.py",
            "*.txt",
            "*.log",
        ],
    )
    print(f"Done! Model available at https://huggingface.co/{HF_REPO_ID}")


if __name__ == "__main__":
    if not os.path.exists(SAFETENSORS_PATH):
        convert()
    else:
        print(f"model.safetensors already exists, skipping conversion.")

    write_config()
    upload()