#!/usr/bin/env python3
"""Step 1: download the base model from Hugging Face into a gitignored cache.

Prints the exact model id, where the weights landed locally, the on-disk size,
and the parameter count. Nothing here is committed: the cache directory is
gitignored, and the reviewer pulls the weights themselves by running this step.
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from demo.common import BASE_MODEL_ID, HF_CACHE, HF_LINK


def _dir_size_bytes(path: Path) -> int:
    return sum(p.stat().st_size for p in path.rglob("*") if p.is_file())


def main() -> int:
    print("=== Step 1: download base model ===\n")
    print(f"  model id:      {BASE_MODEL_ID}")
    print(f"  hugging face:  {HF_LINK}")
    print(f"  cache dir:     {HF_CACHE}  (gitignored)\n")

    from huggingface_hub import snapshot_download
    from transformers import AutoConfig

    local_path = snapshot_download(repo_id=BASE_MODEL_ID)
    config = AutoConfig.from_pretrained(BASE_MODEL_ID)

    n_params = None
    # Derive parameter count cheaply from config where possible, else load.
    try:
        from transformers import AutoModelForCausalLM

        model = AutoModelForCausalLM.from_pretrained(BASE_MODEL_ID)
        n_params = sum(p.numel() for p in model.parameters())
        del model
    except Exception as exc:  # pragma: no cover - informational only
        print(f"  (parameter count skipped: {exc})")

    size_gib = _dir_size_bytes(Path(local_path)) / 2**30
    print(f"  downloaded to: {local_path}")
    print(f"  on-disk size:  {size_gib:.2f} GiB")
    if n_params is not None:
        print(f"  parameters:    {n_params:,} ({n_params / 1e6:.0f}M)")
    print(f"  architecture:  {config.architectures}")
    print(f"  hidden_size={config.hidden_size}, layers={config.num_hidden_layers}")

    print("\nRESULT: PASS - base model present locally and ready to run.")
    return 0


if __name__ == "__main__":
    sys.exit(main())