"""Download public, revision-pinned assets. No login, API inference, or upload."""
import concurrent.futures
import hashlib
import json
from pathlib import Path
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

ROOT = Path(__file__).resolve().parents[1]
MODELS = {
    "bert-mini": ("google/bert_uncased_L-4_H-256_A-4", "387825ce42dbb39b87911cdf8e383ee3b25184f8",
                  ["config.json", "pytorch_model.bin", "vocab.txt", "README.md"]),
    "qwen": ("Qwen/Qwen2.5-1.5B-Instruct", "989aa7980e4cf806f80c7fef2b1adb7bc71aa306",
             ["config.json", "generation_config.json", "tokenizer.json", "tokenizer_config.json",
              "vocab.json", "merges.txt", "model.safetensors", "LICENSE", "README.md"]),
}

def fetch(job):
    url, dest = job
    dest.parent.mkdir(parents=True, exist_ok=True)
    if not dest.exists():
        partial = dest.with_suffix(dest.suffix + ".part")
        session = requests.Session()
        session.mount("https://", HTTPAdapter(max_retries=Retry(total=4, backoff_factor=2, status_forcelist=[429,500,502,503,504])))
        with session.get(url, stream=True, timeout=(30, 180)) as r:
            r.raise_for_status()
            with partial.open("wb") as f:
                for chunk in r.iter_content(2**20):
                    f.write(chunk)
        partial.replace(dest)
    h = hashlib.sha256()
    with dest.open("rb") as f:
        for block in iter(lambda: f.read(2**20), b""):
            h.update(block)
    result = dict(path=str(dest.relative_to(ROOT)), url=url, bytes=dest.stat().st_size, sha256=h.hexdigest())
    print(f"OK {result['path']} ({result['bytes']:,} bytes)", flush=True)
    return result

def main():
    jobs = []
    for name, (repo, rev, files) in MODELS.items():
        jobs += [(f"https://huggingface.co/{repo}/resolve/{rev}/{f}", ROOT / "models" / name / f) for f in files]
    rev = "b1c7086bb997666ef0a2dd393b73f92e824cc487"
    jobs += [(f"https://raw.githubusercontent.com/Krishnarohith10/nlp-getting-started/{rev}/train.csv",
              ROOT / "data/raw/disaster_tweets.csv")]
    with concurrent.futures.ThreadPoolExecutor(4) as pool:
        records = list(pool.map(fetch, jobs))
    (ROOT / "data/asset_manifest.json").write_text(json.dumps(records, ensure_ascii=False, indent=2), encoding="utf-8")

if __name__ == "__main__":
    main()