"""
Download every image a website serves, using its sitemap.

Works on any site with an XML sitemap - Webflow, WordPress, Shopify,
Squarespace, static builds. Handles sitemap index files, keeps the URL
path in the filename so two 'hero.jpg' files can't overwrite each other,
and skips anything already downloaded so you can re-run it safely.

    pip install requests
    python download_images.py https://www.example.com/sitemap.xml
"""

import os
import re
import sys
import xml.etree.ElementTree as ET
from urllib.parse import urlparse, unquote

import requests

OUT = "site-images"
UA = {"User-Agent": "Mozilla/5.0 (compatible; image-backup/1.0)"}
NS = {"ns": "http://www.sitemaps.org/schemas/sitemap/0.9"}
IMG = re.compile(r'https?://[^"\'\s)\\]+?\.(?:jpg|jpeg|png|gif|webp|svg|avif)', re.I)


def fetch(url):
    r = requests.get(url, headers=UA, timeout=20)
    r.raise_for_status()
    return r


def page_urls(sitemap_url, seen_maps=None):
    """Yield every page URL, following sitemap index files one level deep."""
    seen_maps = seen_maps or set()
    if sitemap_url in seen_maps:
        return
    seen_maps.add(sitemap_url)

    root = ET.fromstring(fetch(sitemap_url).content)

    # A sitemap index points at more sitemaps rather than at pages.
    children = [el.text for el in root.findall(".//ns:sitemap/ns:loc", NS)]
    if children:
        print(f"Sitemap index: {len(children)} child sitemaps")
        for child in children:
            yield from page_urls(child, seen_maps)
        return

    for loc in root.findall(".//ns:loc", NS):
        if loc.text:
            yield loc.text.strip()


def safe_name(img_url):
    """Keep the path, not just the basename - two pages often both ship
    a 'hero.jpg', and basename alone silently loses one of them."""
    path = unquote(urlparse(img_url).path).lstrip("/")
    name = re.sub(r"[^A-Za-z0-9._/-]", "_", path).replace("/", "__")
    return name[-180:] or "image"


def main():
    if len(sys.argv) < 2:
        sys.exit("Usage: python download_images.py https://example.com/sitemap.xml")

    os.makedirs(OUT, exist_ok=True)
    seen, saved, failed = set(), 0, 0

    for url in page_urls(sys.argv[1]):
        print(f"Scanning: {url}")
        try:
            html = fetch(url).text
        except Exception as err:
            print(f"  ! could not read page: {err}")
            failed += 1
            continue

        for img_url in IMG.findall(html):
            if img_url in seen:
                continue
            seen.add(img_url)

            dest = os.path.join(OUT, safe_name(img_url))
            if os.path.exists(dest):
                continue

            try:
                data = fetch(img_url).content
            except Exception as err:
                print(f"  ! {img_url}: {err}")
                failed += 1
                continue

            with open(dest, "wb") as fh:
                fh.write(data)
            saved += 1
            print(f"  saved {os.path.basename(dest)}")

    print(f"\nDone. {saved} images in ./{OUT}/ ({len(seen)} found, {failed} failed)")


if __name__ == "__main__":
    main()