"""Instant file scan: inspects a repo's file list and metadata. Downloads no weights.

Safe to run on a free CPU Space: only config.json (plain JSON) is ever fetched.
"""
import json

from huggingface_hub import HfApi, hf_hub_download

# Formats that can execute code when loaded with pickle-based loaders (torch.load, joblib, ...)
PICKLE_EXT = (".bin", ".pt", ".pth", ".ckpt", ".pkl", ".pickle", ".joblib")
DATA_EXT = (".safetensors", ".gguf")
NON_WEIGHT_BIN = ("training_args.bin",)  # pickled, but not weights; still flagged as pickle

GRADES = {
    "A": "Safetensors/GGUF weights only, no custom code required.",
    "B": "Safe weight format, but the model needs trust_remote_code (runs the author's Python).",
    "C": "Safe weights present, plus pickle files that can run code if loaded.",
    "F": "Weights only in pickle format: loading them can execute arbitrary code.",
}


def _base_models(card):
    if not card:
        return []
    bm = getattr(card, "base_model", None) or (card.get("base_model") if isinstance(card, dict) else None)
    if not bm:
        return []
    return [bm] if isinstance(bm, str) else list(bm)


def scan_files(repo_id, revision=None, token=None):
    api = HfApi(token=token)
    try:
        info = api.model_info(repo_id, revision=revision, files_metadata=True, securityStatus=True)
    except TypeError:  # older huggingface_hub without securityStatus
        info = api.model_info(repo_id, revision=revision, files_metadata=True)

    files = [(s.rfilename, s.size or 0) for s in (info.siblings or [])]
    names = [f for f, _ in files]
    pickles = [f for f in names if f.lower().endswith(PICKLE_EXT)]
    safetensors = [f for f in names if f.lower().endswith(".safetensors")]
    gguf = [f for f in names if f.lower().endswith(".gguf")]
    py_files = [f for f in names if f.endswith(".py")]

    auto_map = None
    if "config.json" in names:
        try:
            with open(hf_hub_download(repo_id, "config.json", revision=revision, token=token)) as fh:
                auto_map = json.load(fh).get("auto_map")
        except Exception:
            auto_map = None
    needs_remote_code = bool(auto_map)

    weight_pickles = [f for f in pickles if f.split("/")[-1] not in NON_WEIGHT_BIN]
    if (safetensors or gguf) and not pickles:
        grade = "B" if needs_remote_code else "A"
    elif (safetensors or gguf) and pickles:
        grade = "C"
    elif weight_pickles:
        grade = "F"
    else:
        grade = "B" if needs_remote_code else "A"

    sec = getattr(info, "security_repo_status", None)
    hf_scan = None
    if sec:
        hf_scan = {"scans_done": sec.get("scansDone"),
                   "files_with_issues": [f.get("path") for f in sec.get("filesWithIssues", [])]}

    card = getattr(info, "card_data", None)
    return {
        "repo_id": repo_id,
        "revision": info.sha,
        "grade": grade,
        "grade_meaning": GRADES[grade],
        "pickle_files": pickles,
        "safetensors_files": len(safetensors),
        "gguf_files": len(gguf),
        "custom_code_files": py_files,
        "needs_trust_remote_code": needs_remote_code,
        "declared_base_models": _base_models(card),
        "license": getattr(card, "license", None) if card else None,
        "gated": getattr(info, "gated", None),
        "hf_security_scan": hf_scan,
        "total_size_gb": round(sum(s for _, s in files) / 1e9, 2),
        "deep_scan_supported": bool(safetensors) and not needs_remote_code,
    }