import os, subprocess, zipfile
from pathlib import Path
from huggingface_hub import hf_hub_download, snapshot_download

def writable_data_root():
    preferred = Path(os.environ.get('DATA_ROOT', '/data/amodal'))
    for candidate in (preferred, Path('/tmp/amodal')):
        try:
            candidate.mkdir(parents=True, exist_ok=True)
            test = candidate / '.write_test'
            test.write_text('ok')
            test.unlink(missing_ok=True)
            return candidate
        except Exception:
            continue
    raise RuntimeError('No writable data directory found; tried /data/amodal and /tmp/amodal')

data_root = writable_data_root()
cache_root = data_root.parent/'hf_cache'
os.environ.setdefault('HF_HOME', str(cache_root))
os.environ.setdefault('HUGGINGFACE_HUB_CACHE', str(cache_root/'hub'))
os.environ.setdefault('TRANSFORMERS_CACHE', str(cache_root/'transformers'))
os.environ.setdefault('HOME', str(data_root.parent))
os.environ.setdefault('XDG_CACHE_HOME', str(data_root.parent/'.cache'))
ckpt_root = data_root/'checkpoints'
ckpt_root.mkdir(parents=True, exist_ok=True)
cache_root.mkdir(parents=True, exist_ok=True)
(data_root/'LISAoutput').mkdir(parents=True, exist_ok=True)

def run(cmd, cwd=None, check=True):
    print('+', ' '.join(cmd))
    return subprocess.run(cmd, cwd=cwd, check=check)

# Public direct downloads
files = [
    ('https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth', ckpt_root/'sam_vit_h_4b8939.pth'),
    ('https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth', ckpt_root/'groundingdino_swint_ogc.pth'),
]
for url, dst in files:
    if not dst.exists():
        run(['wget', '-O', str(dst), url])

# HF-hosted RAM++ checkpoint
ram_dst = ckpt_root/'ram_plus_swin_large_14m.pth'
if not ram_dst.exists():
    src = hf_hub_download('xinyu1205/recognize-anything-plus-model', 'ram_plus_swin_large_14m.pth')
    os.symlink(src, ram_dst)

# LISA model: use snapshot cache; app.py can load repo_id directly, but snapshot warms cache.
# Requires HF_TOKEN if the repo/license requires it.
snapshot_download('xinlai/LISA-13B-llama2-v1', local_dir=str(cache_root/'LISA-13B-llama2-v1'), local_dir_use_symlinks=False)

# InstaOrder checkpoint zip is on Google Drive and can be throttled. Try gdown.
insta_ckpt = ckpt_root/'InstaOrder_InstaOrderNet_od.pth.tar'
if not insta_ckpt.exists():
    zip_path = Path('/tmp/InstaOrder_ckpt.zip')
    if not zip_path.exists():
        # Use the raw file id; older gdown versions do not support --fuzzy.
        run(['gdown', '1_GEmCmofLSkJZnidfp4vsQb2Nqq5aqBU', '-O', str(zip_path)], check=False)
    if zip_path.exists():
        try:
            with zipfile.ZipFile(zip_path) as z:
                z.extractall(ckpt_root)
        except zipfile.BadZipFile:
            print('WARNING: downloaded InstaOrder checkpoint is not a valid zip; removing it so a later run can retry.')
            zip_path.unlink(missing_ok=True)
    if not insta_ckpt.exists():
        nested = ckpt_root/'InstaOrder_ckpt'/'InstaOrder_InstaOrderNet_od.pth.tar'
        if nested.exists():
            nested.replace(insta_ckpt)
    if not insta_ckpt.exists():
        print('WARNING: InstaOrder checkpoint missing. Upload InstaOrder_ckpt.zip manually or mirror it to HF.')