#!/usr/bin/env python3
"""Pull RAW rendered maps from every Map Service / Image Service / Tile Service.

For each service:
  - maps_raw/<id>/export_full.png        : /export at full extent, 4096x4096
  - maps_raw/<id>/export_full.json       : /export?f=json (image href + extent)
  - maps_raw/<id>/_service.json          : MapServer/ImageServer metadata
  - maps_raw/<id>/legend.json            : /legend (symbology)
  - maps_raw/<id>/tiles/{z}/{y}/{x}.png  : tile pyramid (z0-z6 only) if cached
"""
import json, time, threading, urllib.request, urllib.error
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

ROOT = Path(__file__).parent
INDEX = ROOT / "items_index.json"
OUT   = ROOT / "maps_raw"
LOG   = ROOT / "logs" / "maps_raw.log"
STAT  = ROOT / "logs" / "maps_raw_status.json"
UA    = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36"
WORKERS = 6
TIMEOUT = 90
MAX_TILE_Z = 6   # cap tile pyramid depth to keep it manageable
EXPORT_SIZE = "4096,4096"

_lock = threading.Lock()
_stats = {"services":0, "exports":0, "tiles":0, "skip":0, "err":0, "start":time.time()}

def log(msg):
    line = f"[{time.strftime('%H:%M:%S')}] {msg}"
    with _lock:
        print(line, flush=True)
        with open(LOG, "a", encoding="utf-8") as f:
            f.write(line+"\n")

def fetch(url, binary=False, retries=2):
    for i in range(retries+1):
        try:
            req = urllib.request.Request(url, headers={"User-Agent": UA})
            with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
                return r.read() if binary else r.read().decode("utf-8", errors="replace")
        except urllib.error.HTTPError as e:
            if e.code in (404,401,403): return None
            if i==retries: return None
        except Exception:
            if i==retries: return None
        time.sleep(1+i)
    return None

def dump_service(it):
    iid = it["id"]; url = it.get("url","").rstrip("/")
    typ = it.get("type","")
    if not url or typ not in {"Map Service","Image Service"}: return
    d = OUT / iid
    if (d / "_done.txt").exists():
        with _lock: _stats["skip"] += 1
        return
    d.mkdir(parents=True, exist_ok=True)

    # service metadata
    meta_txt = fetch(f"{url}?f=json")
    if not meta_txt:
        with _lock: _stats["err"] += 1
        return
    (d / "_service.json").write_text(meta_txt, encoding="utf-8")
    try: meta = json.loads(meta_txt)
    except: meta = {}
    with _lock: _stats["services"] += 1

    # legend
    lg = fetch(f"{url}/legend?f=json")
    if lg: (d / "legend.json").write_text(lg, encoding="utf-8")

    # full-extent export
    ext = meta.get("fullExtent") or meta.get("initialExtent") or meta.get("extent")
    if ext:
        bbox = f"{ext['xmin']},{ext['ymin']},{ext['xmax']},{ext['ymax']}"
        sr   = ext.get("spatialReference",{}).get("wkid", 102100)
        eurl = (f"{url}/export?bbox={bbox}&bboxSR={sr}&imageSR={sr}"
                f"&size={EXPORT_SIZE}&format=png32&transparent=true&f=image")
        png = fetch(eurl, binary=True)
        if png and len(png) > 1000:
            (d / "export_full.png").write_bytes(png)
            with _lock: _stats["exports"] += 1
        # also save the JSON form (gives image href + adjusted extent)
        jurl = eurl.rsplit("&f=",1)[0] + "&f=json"
        ej = fetch(jurl)
        if ej: (d / "export_full.json").write_text(ej, encoding="utf-8")

    # tile pyramid if cached
    if meta.get("singleFusedMapCache"):
        lods = meta.get("tileInfo",{}).get("lods",[])
        for lod in lods:
            z = lod.get("level")
            if z is None or z > MAX_TILE_Z: continue
            # use tilemap to know which tiles exist (1 call per z if available)
            tm_url = f"{url}/tilemap/{z}/0/0/{1<<z}/{1<<z}"
            tm = fetch(tm_url)
            present = None
            if tm:
                try:
                    tmd = json.loads(tm)
                    present = tmd.get("data")
                except: pass
            if present:
                cols = 1<<z
                idx = 0
                for row in range(cols):
                    for col in range(cols):
                        if idx < len(present) and present[idx]:
                            out = d / "tiles" / str(z) / str(row) / f"{col}.png"
                            if not out.exists():
                                out.parent.mkdir(parents=True, exist_ok=True)
                                tile = fetch(f"{url}/tile/{z}/{row}/{col}", binary=True)
                                if tile:
                                    out.write_bytes(tile)
                                    with _lock: _stats["tiles"] += 1
                        idx += 1
            else:
                # no tilemap — try 0/0/0 to see if any tiles at all
                if z == 0:
                    t = fetch(f"{url}/tile/0/0/0", binary=True)
                    if t:
                        out = d / "tiles" / "0" / "0" / "0.png"
                        out.parent.mkdir(parents=True, exist_ok=True)
                        out.write_bytes(t)
                        with _lock: _stats["tiles"] += 1
    (d / "_done.txt").write_text("ok")

def status_writer():
    while True:
        time.sleep(15)
        with _lock: s = dict(_stats)
        s["elapsed"] = int(time.time()-s["start"])
        STAT.write_text(json.dumps(s, indent=2))

def main():
    OUT.mkdir(exist_ok=True)
    LOG.parent.mkdir(exist_ok=True)
    items = json.load(open(INDEX, encoding="utf-8"))
    services = [i for i in items if i.get("type") in {"Map Service","Image Service"}]
    log(f"{len(services)} map/image services to render")
    threading.Thread(target=status_writer, daemon=True).start()
    with ThreadPoolExecutor(max_workers=WORKERS) as ex:
        futs = [ex.submit(dump_service, s) for s in services]
        done = 0
        for f in as_completed(futs):
            done += 1
            if done % 10 == 0:
                with _lock: s = dict(_stats)
                log(f"progress {done}/{len(services)} svc={s['services']} png={s['exports']} tiles={s['tiles']} err={s['err']}")
    log(f"DONE {_stats}")

if __name__ == "__main__":
    main()