scripts/download_met_museum.py

"""
The Metropolitan Museum of Art Open Access downloader.

The Met has 492,000+ CC0 images, with strong Mughal miniature and
South Asian art holdings. Their API is fast, free, requires no key,
and doesn't block cloud IPs.

API docs: https://metmuseum.github.io/

Output: assets/datasets/raw/<style>/<style>_NNN_<title>.jpg

Usage:
    python scripts/download_met_museum.py --all
    python scripts/download_met_museum.py --style mughal --count 60
"""
from __future__ import annotations

import argparse
import json
import logging
import re
import time
from pathlib import Path
from typing import List

import requests
from PIL import Image
from io import BytesIO

log = logging.getLogger(__name__)

# The Met API — completely open, no key required
MET_API = "https://collectionapi.metmuseum.org/public/collection/v1"

HEADERS = {
    "User-Agent": "IndicHeritageStudio/2.0 (educational hackathon project)",
}

# Search queries per style — tuned for The Met's collection
# The Met has particularly strong Mughal holdings
STYLE_QUERIES = {
    "madhubani": [
        "madhubani",
        "mithila",
    ],
    "warli": [
        "warli",
    ],
    "pattachitra": [
        "pattachitra",
        "orissa painting",
    ],
    "mughal": [
        "mughal painting",
        "mughal miniature",
        "mughal manuscript",
        "imperial mughal",
        "akbar manuscript",
        "jahangir painting",
        "shah jahan painting",
        "mughal album page",
        "mughal portraiture",
        "hamzanama",
    ],
    "tanjore": [
        "tanjore painting",
        "thanjavur",
        "south indian painting",
    ],
}


def search_met(query: str) -> List[int]:
    """Search The Met API for object IDs matching the query."""
    try:
        resp = requests.get(
            f"{MET_API}/search",
            params={"q": query, "hasImages": "true", "medium": "Paintings"},
            headers=HEADERS,
            timeout=30,
        )
        resp.raise_for_status()
        data = resp.json()
        return data.get("objectIDs", []) or []
    except Exception as exc:
        log.warning(f"Met search failed for '{query}': {exc}")
        return []


def get_object(object_id: int) -> dict:
    """Fetch full metadata for a Met object."""
    try:
        resp = requests.get(
            f"{MET_API}/objects/{object_id}",
            headers=HEADERS,
            timeout=30,
        )
        resp.raise_for_status()
        return resp.json()
    except Exception as exc:
        log.debug(f"Met object {object_id} fetch failed: {exc}")
        return {}


def download_image(url: str, out_path: Path, target_size: int = 1024) -> bool:
    """Download and resize an image."""
    try:
        resp = requests.get(url, headers=HEADERS, timeout=60, allow_redirects=True)
        if resp.status_code != 200 or len(resp.content) < 1000:
            return False
        img = Image.open(BytesIO(resp.content)).convert("RGB")
        w, h = img.size
        if w >= h:
            new_w, new_h = target_size, int(h * target_size / w)
        else:
            new_w, new_h = int(w * target_size / h), target_size
        img = img.resize((new_w, new_h), Image.LANCZOS)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        img.save(out_path, "JPEG", quality=95)
        return True
    except Exception as exc:
        log.debug(f"Download failed {url}: {exc}")
        return False


def download_style(style_id: str, total_count: int = 40) -> int:
    """Download `total_count` images for one style."""
    queries = STYLE_QUERIES.get(style_id, [])
    if not queries:
        log.error(f"No queries for style '{style_id}'")
        return 0

    out_dir = Path(f"assets/datasets/raw/{style_id}")
    out_dir.mkdir(parents=True, exist_ok=True)

    downloaded = 0
    seen_urls = set()
    seen_ids = set()

    for query in queries:
        if downloaded >= total_count:
            break
        log.info(f"[{style_id}] Searching Met Museum: '{query}'")
        object_ids = search_met(query)
        log.info(f"[{style_id}] Found {len(object_ids)} candidate objects")

        for oid in object_ids:
            if downloaded >= total_count:
                break
            if oid in seen_ids:
                continue
            seen_ids.add(oid)

            obj = get_object(oid)
            if not obj:
                continue

            # Get the primary image URL
            img_url = obj.get("primaryImage") or obj.get("primaryImageSmall")
            if not img_url:
                continue
            if img_url in seen_urls:
                continue
            seen_urls.add(img_url)

            # Verify it's a painting and matches our style
            title = obj.get("title", "")[:80] or "untitled"
            classification = obj.get("classification", "")
            country = obj.get("country", "")
            culture = obj.get("culture", "")
            department = obj.get("department", "")

            # Filter: must be in Asian Art or Paintings department
            if department and department not in ("Asian Art", "Paintings"):
                continue

            # Skip if obviously not relevant
            obj_text = f"{title} {classification} {country} {culture}".lower()
            if style_id == "mughal" and "mughal" not in obj_text and "india" not in obj_text:
                continue
            if style_id == "tanjore" and "tanjore" not in obj_text and "tamil" not in obj_text and "south india" not in obj_text:
                continue

            safe_title = re.sub(r"[^a-zA-Z0-9_-]", "_", title)[:50]
            out_path = out_dir / f"{style_id}_{downloaded:03d}_{safe_title}.jpg"

            if download_image(img_url, out_path):
                downloaded += 1
                log.info(f"  ✓ [{downloaded}/{total_count}] {out_path.name} "
                        f"(dept={department}, culture={culture[:30]})")
                time.sleep(0.3)

    log.info(f"[{style_id}] Downloaded {downloaded}/{total_count} from Met Museum")
    return downloaded


def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")

    p = argparse.ArgumentParser(description="Download heritage art from The Met Museum Open Access")
    p.add_argument("--style", choices=list(STYLE_QUERIES.keys()))
    p.add_argument("--all", action="store_true")
    p.add_argument("--count", type=int, default=40)
    args = p.parse_args()

    styles = [args.style] if args.style else list(STYLE_QUERIES.keys())

    log.info("=== The Met Museum Open Access Downloader ===")
    log.info(f"Styles: {styles}")
    log.info(f"Per style: {args.count}")

    total = 0
    for style in styles:
        log.info(f"\n--- {style.upper()} ---")
        n = download_style(style, total_count=args.count)
        total += n

    log.info(f"\n=== Done. Total: {total} images ===")


if __name__ == "__main__":
    main()