#!/usr/bin/env python3
"""Download sitemaps and write URL lists. Respects fetch.py rate limits."""
import sys
from pathlib import Path
from xml.etree import ElementTree as ET

sys.path.insert(0, str(Path(__file__).parent))
from fetch import get

ROOT = Path("/Users/hattr/Documents/wa-land-developers/raw")
NS = {"sm": "http://www.sitemaps.org/schemas/sitemap/0.9"}


def locs(xml_bytes):
    # strip default ns for easier parse if needed
    try:
        root = ET.fromstring(xml_bytes)
    except ET.ParseError as e:
        print("PARSE FAIL", e, xml_bytes[:200])
        return [], []
    sitemaps, urls = [], []
    for el in root.iter():
        tag = el.tag.split("}")[-1]
        if tag == "loc" and el.text:
            parent = None
    # walk properly
    for child in list(root):
        tag = child.tag.split("}")[-1]
        loc = None
        for sub in child.iter():
            if sub.tag.split("}")[-1] == "loc" and sub.text:
                loc = sub.text.strip()
                break
        if not loc:
            continue
        if tag == "sitemap":
            sitemaps.append(loc)
        elif tag == "url":
            urls.append(loc)
    return sitemaps, urls


def pull(name, index_url):
    print("INDEX", name, index_url, flush=True)
    status, final, content, cached = get(index_url, timeout=60)
    (ROOT / f"{name}-index.xml").write_bytes(content)
    print(f"  status={status} cached={cached} bytes={len(content)}", flush=True)
    sitemaps, urls = locs(content)
    all_urls = list(urls)
    for sm in sitemaps:
        print("  CHILD", sm, flush=True)
        st, fin, body, c = get(sm, timeout=60)
        safe = sm.split("/")[-1].split("?")[0] or "child"
        (ROOT / f"{name}--{safe}").write_bytes(body)
        _, child_urls = locs(body)
        print(f"    status={st} urls={len(child_urls)} cached={c}", flush=True)
        all_urls.extend(child_urls)
    out = ROOT / f"{name}-urls.txt"
    out.write_text("\n".join(all_urls) + ("\n" if all_urls else ""))
    print(f"WROTE {out} count={len(all_urls)}", flush=True)


if __name__ == "__main__":
    jobs = [
        ("udia", "https://www.udiawa.com.au/sitemap_index.xml"),
        ("devwa", "https://developmentwa.com.au/index.php?option=com_jmap&view=sitemap&format=xml"),
        ("cw", "https://www.cossillwebley.com.au/sitemap_index.xml"),
        ("upd", "https://www.upd.com.au/sitemap_index.xml"),
    ]
    only = sys.argv[1:] 
    for name, url in jobs:
        if only and name not in only:
            continue
        try:
            pull(name, url)
        except Exception as e:
            print("FAIL", name, type(e).__name__, e, flush=True)