scripts/download_smithsonian.py

"""
Smithsonian Open Access downloader.

The Smithsonian Institution has 4.5 million CC0 (public domain) images,
including a large Asian art collection. Their API is fast, free, and
doesn't block cloud IPs (unlike Wikimedia).

API docs: https://edan.si.edu/openaccess/apidocs/

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

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

import argparse
import json
import logging
import re
import time
from pathlib import Path
from typing import List
from urllib.parse import quote

import requests
from PIL import Image
from io import BytesIO

log = logging.getLogger(__name__)

# Smithsonian API — no API key required for basic search
API_BASE = "https://api.si.edu/openaccess/api/v1.0/content"
API_KEY = "Demo"  # Smithsonian allows demo key for low-volume use

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

# Search queries per style — tuned for Smithsonian's collection
# Smithsonian uses their own subject keywords; these are validated to return results.
STYLE_QUERIES = {
    "madhubani": [
        "madhubani",
        "mithila painting",
        "indian folk painting bihar",
    ],
    "warli": [
        "warli",
        "warli painting",
        "indian tribal art maharashtra",
    ],
    "pattachitra": [
        "pattachitra",
        "odisha painting",
        "orissa scroll painting",
    ],
    "mughal": [
        "mughal painting",
        "mughal miniature",
        "mughal manuscript",
        "indian miniature painting",
        "akbar painting",
        "jahangir painting",
    ],
    "tanjore": [
        "tanjore painting",
        "thanjavur painting",
        "south indian painting",
        "tamil painting",
    ],
}


def search_smithsonian(query: str, rows: int = 50) -> List[dict]:
    """Search Smithsonian Open Access for images matching the query."""
    results = []
    params = {
        "api_key": API_KEY,
        "q": query,
        "fq": 'type:"emuseum-images"',  # only return image records
        "rows": rows,
        "start": 0,
    }

    try:
        resp = requests.get(API_BASE, params=params, headers=HEADERS, timeout=30)
        resp.raise_for_status()
        data = resp.json()
    except Exception as exc:
        log.warning(f"Smithsonian search failed for '{query}': {exc}")
        return []

    rows = data.get("response", {}).get("rows", [])
    for row in rows:
        # Extract image URL + metadata from Smithsonian's nested JSON
        content = row.get("content", {})
        descriptiveNonrepeating = content.get("descriptiveNonrepeating", {})
        online_media = descriptiveNonrepeating.get("online_media", {})
        media_list = online_media.get("media", []) if isinstance(online_media, dict) else []

        if not media_list:
            continue

        # Get title
        title = (content.get("freetext", {}).get("title", {}) or {}).get("label", "")[:80]
        if not title:
            title = (row.get("title", "") or "")[:80]

        # Get first image URL
        for media in media_list:
            if not isinstance(media, dict):
                continue
            media_type = media.get("type", "")
            if media_type not in ("Images", "Image"):
                continue
            content_url = media.get("content", "")
            if not content_url:
                continue
            thumbnail = media.get("thumbnail", "")
            usage = media.get("usage", {}).get("access", "")
            results.append({
                "title": title or "untitled",
                "url": content_url,
                "thumbnail": thumbnail,
                "usage": usage,
                "id": row.get("id", ""),
            })
            break  # only take first image per record

    return results


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 across multiple queries."""
    queries = STYLE_QUERIES.get(style_id, [])
    if not queries:
        log.error(f"No queries for style '{style_id}'")
        return 0

    per_query = max(10, total_count // len(queries))
    out_dir = Path(f"assets/datasets/raw/{style_id}")
    out_dir.mkdir(parents=True, exist_ok=True)

    downloaded = 0
    seen_urls = set()

    for query in queries:
        if downloaded >= total_count:
            break
        log.info(f"[{style_id}] Searching Smithsonian: '{query}'")
        results = search_smithsonian(query, rows=per_query)
        log.info(f"[{style_id}] Found {len(results)} candidates")

        for r in results:
            if downloaded >= total_count:
                break
            if r["url"] in seen_urls:
                continue
            seen_urls.add(r["url"])

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

            if download_image(r["url"], out_path):
                downloaded += 1
                log.info(f"  ✓ [{downloaded}/{total_count}] {out_path.name} ({r['usage']})")
                time.sleep(0.3)
            else:
                log.debug(f"  ✗ {r['title']}")

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


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

    p = argparse.ArgumentParser(description="Download heritage art from Smithsonian 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("=== Smithsonian 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()