"""
setup_local_model.py
────────────────────
Run this ONCE to bundle the HuggingFace config files alongside
the local pytorch_model.bin you already downloaded.

Usage:
    python setup_local_model.py
"""

import os
import shutil
import json
import urllib.request

# ── Paths ───────────────────────────────────────────────────
BIN_SRC  = r"C:\Users\mani8\Downloads\pytorch_model.bin"
MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "local_model")

# ── Files to fetch from HuggingFace (tiny JSONs, < 5 KB each) ──
HF_BASE = "https://huggingface.co/umm-maybe/AI-image-detector/resolve/main"
CONFIG_FILES = [
    "config.json",
    "preprocessor_config.json",
]

# ────────────────────────────────────────────────────────────
def main():
    os.makedirs(MODEL_DIR, exist_ok=True)
    print(f"📁 Model directory: {MODEL_DIR}\n")

    # 1. Copy the .bin weights
    bin_dst = os.path.join(MODEL_DIR, "pytorch_model.bin")
    if os.path.exists(bin_dst):
        print(f"✅ pytorch_model.bin already present — skipping copy")
    else:
        if not os.path.exists(BIN_SRC):
            raise FileNotFoundError(f"pytorch_model.bin not found at:\n  {BIN_SRC}")
        print(f"📋 Copying pytorch_model.bin ... ", end="", flush=True)
        shutil.copy2(BIN_SRC, bin_dst)
        size_mb = os.path.getsize(bin_dst) / 1_048_576
        print(f"done  ({size_mb:.1f} MB)")

    # 2. Download config files
    for fname in CONFIG_FILES:
        dst = os.path.join(MODEL_DIR, fname)
        if os.path.exists(dst):
            print(f"✅ {fname} already present — skipping")
            continue
        url = f"{HF_BASE}/{fname}"
        print(f"🌐 Downloading {fname} ...", end="", flush=True)
        try:
            urllib.request.urlretrieve(url, dst)
            print(" done")
        except Exception as e:
            print(f"\n⚠ Could not download {fname}: {e}")
            print("   Writing fallback config...")
            _write_fallback(fname, dst)

    # 3. Verify
    print("\n── Contents of local_model/ ─────────────────────")
    for f in os.listdir(MODEL_DIR):
        size = os.path.getsize(os.path.join(MODEL_DIR, f))
        print(f"   {f:<35} {size/1024:>8.1f} KB")

    print("\n✅ Local model ready!")
    print("   You can now run:  python app.py")


# ── Fallback hardcoded configs (from umm-maybe/AI-image-detector) ──
def _write_fallback(fname, dst):
    if fname == "config.json":
        cfg = {
            "architectures": ["ViTForImageClassification"],
            "hidden_size": 768,
            "id2label": {"0": "artificial", "1": "real"},
            "label2id": {"artificial": 0, "real": 1},
            "model_type": "vit",
            "num_attention_heads": 12,
            "num_channels": 3,
            "num_hidden_layers": 12,
            "num_labels": 2,
            "patch_size": 16,
            "image_size": 224,
            "torch_dtype": "float32"
        }
    elif fname == "preprocessor_config.json":
        cfg = {
            "do_normalize": True,
            "do_rescale": True,
            "do_resize": True,
            "feature_extractor_type": "ViTFeatureExtractor",
            "image_mean": [0.5, 0.5, 0.5],
            "image_std":  [0.5, 0.5, 0.5],
            "resample": 2,
            "rescale_factor": 0.00392156862745098,
            "size": {"height": 224, "width": 224}
        }
    else:
        cfg = {}

    with open(dst, "w") as f:
        json.dump(cfg, f, indent=2)
    print(f"   Fallback {fname} written.")


if __name__ == "__main__":
    main()