#!/usr/bin/env python3
"""For every Web Map / Dashboard / StoryMap / Web Experience / Web Mapping App
JSON already saved under data/, find every referenced service URL and item ID,
then dump the underlying layer data (features + metadata + legend).

Output:
  referenced/<sha>/{_service.json, _legend.json, L<k>_meta.json, L<k>_pN.geojson}
  referenced/_url_index.json  (sha -> {url, source_items[]})
  referenced/_item_ref_index.json  (referenced item id -> {source_items[]})
"""
import json, re, hashlib, time, threading, urllib.request, urllib.error, urllib.parse
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

ROOT = Path(__file__).parent
DATA = ROOT / "data"
OUT  = ROOT / "referenced"
LOG  = ROOT / "logs" / "referenced.log"
STAT = ROOT / "logs" / "referenced_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 = 8
TIMEOUT = 60
PAGE = 5000

_lock = threading.Lock()
_stats = {"jsons_parsed":0, "urls_found":0, "urls_dumped":0, "layers":0, "err":0, "start":time.time()}
_seen = set()

URL_RE = re.compile(r'https?://[^\s"\'<>)]+')
SERVICE_TAIL = re.compile(r"(.+?/(?:MapServer|FeatureServer|ImageServer|SceneServer|VectorTileServer|WFSServer|WMSServer|StreamServer))(?:/\d+)?/?$", re.I)

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 sha(s): return hashlib.sha1(s.encode()).hexdigest()[:16]

def normalize_service_url(u):
    """Strip /<layer-id> off the end so we hit the service root."""
    m = SERVICE_TAIL.match(u.rstrip("/"))
    return m.group(1) if m else None

def walk_json(obj, urls, items):
    """Recurse, collecting service URLs and itemIds."""
    if isinstance(obj, dict):
        for k,v in obj.items():
            if k=="url" and isinstance(v,str):
                base = normalize_service_url(v)
                if base: urls.add(base)
            elif k in ("itemId","webmap","portalItemID","resourceId") and isinstance(v,str) and len(v)==32:
                items.add(v)
            walk_json(v, urls, items)
    elif isinstance(obj, list):
        for v in obj: walk_json(v, urls, items)
    elif isinstance(obj, str):
        # regex sweep for any embedded service URLs
        for m in URL_RE.findall(obj):
            b = normalize_service_url(m)
            if b: urls.add(b)

def dump_service(url, sources):
    h = sha(url)
    d = OUT / h
    if (d / "_done.txt").exists():
        return
    d.mkdir(parents=True, exist_ok=True)
    (d / "_url.txt").write_text(url + "\n" + "\n".join(sorted(sources)), encoding="utf-8")
    meta_txt = fetch(f"{url}?f=json")
    if not meta_txt:
        with _lock: _stats["err"] += 1
        (d / "_failed.txt").write_text("metadata fetch failed")
        return
    (d / "_service.json").write_text(meta_txt, encoding="utf-8")
    try: meta = json.loads(meta_txt)
    except: meta = {}
    if isinstance(meta, dict) and meta.get("error"):
        (d / "_failed.txt").write_text(json.dumps(meta.get("error")))
        return
    lg = fetch(f"{url}/legend?f=json")
    if lg: (d / "_legend.json").write_text(lg, encoding="utf-8")
    layers = (meta.get("layers") or []) + (meta.get("tables") or [])
    # ImageServers don't have layers — try /exportImage and metadata only
    if not layers and ("/ImageServer" in url or "/StreamServer" in url):
        ext = meta.get("fullExtent") or meta.get("extent")
        if ext:
            bbox = f"{ext['xmin']},{ext['ymin']},{ext['xmax']},{ext['ymax']}"
            sr   = ext.get("spatialReference",{}).get("wkid",4326)
            iurl = f"{url}/exportImage?bbox={bbox}&bboxSR={sr}&imageSR={sr}&size=2048,2048&format=png32&f=image"
            img = fetch(iurl, binary=True)
            if img and len(img) > 1000:
                (d / "exportImage.png").write_bytes(img)
        (d / "_done.txt").write_text("ok")
        return
    for L in layers:
        lid = L.get("id")
        lmeta = fetch(f"{url}/{lid}?f=json")
        if lmeta: (d / f"L{lid}_meta.json").write_text(lmeta, encoding="utf-8")
        offset = 0; page = 0
        while True:
            q = (f"{url}/{lid}/query?where=1%3D1&outFields=*&returnGeometry=true"
                 f"&f=geojson&resultOffset={offset}&resultRecordCount={PAGE}")
            qr = fetch(q)
            if not qr: break
            try: qd = json.loads(qr)
            except: break
            feats = qd.get("features", [])
            if not feats and offset > 0: break
            (d / f"L{lid}_p{page}.geojson").write_text(qr, encoding="utf-8")
            with _lock: _stats["layers"] += 1
            if len(feats) < PAGE: break
            offset += PAGE; page += 1
            if page > 200: break
    (d / "_done.txt").write_text("ok")
    with _lock: _stats["urls_dumped"] += 1

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)
    threading.Thread(target=status_writer, daemon=True).start()

    # 1. Walk every .json file in data/ collecting service URLs and item refs
    url_sources = {}  # url -> set of source item IDs
    item_refs   = {}  # itemId -> set of source item IDs
    for jp in DATA.glob("*.json"):
        try: txt = jp.read_text(encoding="utf-8")
        except: continue
        try: obj = json.loads(txt)
        except: continue
        urls = set(); items = set()
        walk_json(obj, urls, items)
        src_id = jp.stem
        for u in urls:
            url_sources.setdefault(u, set()).add(src_id)
        for i in items:
            item_refs.setdefault(i, set()).add(src_id)
        with _lock: _stats["jsons_parsed"] += 1

    log(f"parsed {_stats['jsons_parsed']} json files")
    log(f"found {len(url_sources)} unique service URLs and {len(item_refs)} item-id refs")
    with _lock: _stats["urls_found"] = len(url_sources)

    # write the indexes
    (OUT / "_url_index.json").write_text(json.dumps(
        {u: sorted(list(s)) for u,s in url_sources.items()}, indent=2), encoding="utf-8")
    (OUT / "_item_ref_index.json").write_text(json.dumps(
        {i: sorted(list(s)) for i,s in item_refs.items()}, indent=2), encoding="utf-8")

    # 2. Dump each unique service
    with ThreadPoolExecutor(max_workers=WORKERS) as ex:
        futs = {ex.submit(dump_service, u, list(s)): u for u,s in url_sources.items()}
        done = 0
        for f in as_completed(futs):
            done += 1
            if done % 25 == 0:
                with _lock: s = dict(_stats)
                log(f"progress {done}/{len(futs)} dumped={s['urls_dumped']} layers={s['layers']} err={s['err']}")

    log(f"DONE {_stats}")

if __name__ == "__main__":
    main()