#!/bin/sh
# Dump EVERY cached object of one or more domains into ./<domain>/... preserving
# the real URL paths and filenames. Pass --images to take only images.
#
# Usage:   ./dump-images.sh <domain> [domain ...] [--images] [--out DIR] [-j N]
# Env:     NSIN_API_KEY (required)  NSIN_API  OUTDIR  IMAGES=1  CONCURRENCY  INSECURE=1
#
# Talks to the Nsin REST API with an `nsin_…` API key — no Redis/Dragonfly
# credentials, no direct L2 access.
#
# Two steps, because the cache-key registry stores metadata, NOT bodies:
#   1. GET /domains/<name>/cache/keys  (paged)  -> every live cached URL
#   2. HTTPS GET each URL              (edge)   -> the bytes, written to disk
# Step 2 hits the CDN edge, so an entry still in cache is served from cache.
#
# Needs only python3 (stdlib) — no pip, no curl.

set -eu

NSIN_API="${NSIN_API:-https://api.nsin.ir}"
NSIN_API_KEY="${NSIN_API_KEY:-}"
OUTDIR="${OUTDIR:-$PWD}"
IMAGES="${IMAGES:-}"
CONCURRENCY="${CONCURRENCY:-8}"
INSECURE="${INSECURE:-}"

DOMAINS=""
while [ $# -gt 0 ]; do
	case "$1" in
	--images | -i) IMAGES=1 ;;
	--all | -a) ;; # everything is the default now; kept so old invocations still work
	--out | -o)
		shift
		OUTDIR="${1:?--out needs a directory}"
		;;
	-j | --concurrency)
		shift
		CONCURRENCY="${1:?-j needs a number}"
		;;
	-k | --insecure) INSECURE=1 ;;
	-h | --help)
		sed -n '2,/^$/p' "$0" | sed 's/^# \{0,1\}//'
		exit 0
		;;
	-*)
		echo "unknown flag: $1" >&2
		exit 2
		;;
	*) DOMAINS="$DOMAINS $1" ;;
	esac
	shift
done

# One domain is required — everything else has a sane default.
if [ -z "${DOMAINS# }" ]; then
	echo "usage: $0 <domain> [domain ...] [--images] [--out DIR] [-j N]" >&2
	exit 2
fi
if [ -z "$NSIN_API_KEY" ]; then
	echo "NSIN_API_KEY is not set — create one in the Nsin panel (Settings → API keys)." >&2
	exit 2
fi

export NSIN_API NSIN_API_KEY OUTDIR IMAGES CONCURRENCY INSECURE DOMAINS

exec python3 - <<'PY'
import hashlib, json, os, re, ssl, sys, threading, time
from concurrent.futures import ThreadPoolExecutor
from urllib.parse import quote, unquote, urlencode
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError

API = os.environ["NSIN_API"].rstrip("/")
KEY = os.environ["NSIN_API_KEY"]
OUTDIR = os.environ["OUTDIR"]
IMAGES = os.environ.get("IMAGES") == "1"  # off = take everything that's cached
DOMAINS = os.environ["DOMAINS"].split()
WORKERS = max(1, min(32, int(os.environ.get("CONCURRENCY") or 8)))
CTX = ssl._create_unverified_context() if os.environ.get("INSECURE") == "1" else None

PAGE = 500  # cachereg.ListMaxLimit — anything larger is silently clamped to 100
IMG_EXT = re.compile(r"\.(jpe?g|png|gif|webp|avif|svg|bmp|ico|tiff?|heic)$", re.I)
# Only consulted for URLs that carry no extension of their own (image proxies,
# /api/… JSON, pretty URLs) — so the file on disk is openable by double-click.
EXT_FOR_CTYPE = {
    "image/jpeg": ".jpg", "image/png": ".png", "image/gif": ".gif",
    "image/webp": ".webp", "image/avif": ".avif", "image/svg+xml": ".svg",
    "image/x-icon": ".ico", "image/vnd.microsoft.icon": ".ico",
    "image/bmp": ".bmp", "image/tiff": ".tiff", "image/heic": ".heic",
    "text/html": ".html", "text/css": ".css", "text/plain": ".txt",
    "text/javascript": ".js", "application/javascript": ".js",
    "application/json": ".json", "application/xml": ".xml", "text/xml": ".xml",
    "application/pdf": ".pdf", "application/zip": ".zip",
    "font/woff2": ".woff2", "font/woff": ".woff", "font/ttf": ".ttf",
    "video/mp4": ".mp4", "audio/mpeg": ".mp3",
}

# ---------- terminal ----------

TTY = sys.stderr.isatty() and not os.environ.get("NO_COLOR")


def c(code, s):
    return "\x1b[%sm%s\x1b[0m" % (code, s) if TTY else s


dim = lambda s: c("2", s)
bold = lambda s: c("1", s)
green = lambda s: c("32", s)
red = lambda s: c("31", s)
yellow = lambda s: c("33", s)
cyan = lambda s: c("36", s)


def human(n):
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024 or unit == "GB":
            return "%.0f %s" % (n, unit) if unit == "B" else "%.1f %s" % (n, unit)
        n /= 1024.0


def clock(sec):
    if sec is None or sec != sec or sec in (float("inf"),):
        return "--:--"
    sec = int(sec)
    return "%02d:%02d" % (sec // 60, sec % 60)


class Progress:
    """A live multi-line status block redrawn in place.

    Everything that wants to print goes through log(): it wipes the block,
    writes the line, then redraws — so finished-file lines scroll normally
    above a bar that never gets torn in half.
    """

    BLOCKS = "▏▎▍▌▋▊▉█"

    def __init__(self):
        self.lock = threading.RLock()
        self.lines = 0
        self.rows = []
        self.last = 0.0
        if TTY:
            sys.stderr.write("\x1b[?25l")  # hide cursor

    def _wipe(self):
        if TTY and self.lines:
            sys.stderr.write("\x1b[%dF\x1b[0J" % self.lines)
            self.lines = 0

    def _draw(self):
        if not TTY or not self.rows:
            return
        sys.stderr.write("\n".join(self.rows) + "\n")
        self.lines = len(self.rows)
        sys.stderr.flush()

    def set(self, rows, force=False):
        with self.lock:
            now = time.monotonic()
            if not force and now - self.last < 0.06:  # ~16 fps ceiling
                self.rows = rows
                return
            self.last = now
            self._wipe()
            self.rows = rows
            self._draw()

    def log(self, msg):
        with self.lock:
            self._wipe()
            sys.stderr.write(msg + "\n")
            self._draw()
            sys.stderr.flush()

    def close(self):
        with self.lock:
            self._wipe()
            self.rows = []
            if TTY:
                sys.stderr.write("\x1b[?25h")  # restore cursor
            sys.stderr.flush()

    @classmethod
    def bar(cls, frac, width=28):
        frac = 0.0 if frac < 0 else 1.0 if frac > 1 else frac
        full = int(frac * width)
        rem = int((frac * width - full) * 8)
        cells = "█" * full
        if full < width:
            cells += (cls.BLOCKS[rem - 1] if rem else " ")
            cells += " " * (width - full - 1)
        return c("36", cells) if TTY else cells


P = Progress()

# ---------- API ----------


def api_get(path, params):
    url = "%s%s?%s" % (API, path, urlencode(params))
    req = Request(url, headers={
        "Authorization": "Bearer " + KEY,
        "Accept": "application/json",
        "User-Agent": "nsin-dump-images/2",
    })
    with urlopen(req, timeout=60, context=CTX) as r:
        return json.loads(r.read() or b"{}")


def list_keys(domain):
    """Page the whole cache-key registry for one domain.

    Rows are deduped on the display identity (host+path+query): the same URL is
    cached independently on every edge node and once per variant (|a=webp,
    |d=mobile, …), and all of those collapse to one file on disk.
    """
    rows, seen, offset, total = [], set(), 0, None
    while True:
        try:
            page = api_get("/domains/%s/cache/keys" % quote(domain, safe=""),
                           {"limit": PAGE, "offset": offset, "sort": "size", "dir": "desc"})
        except HTTPError as e:
            body = (e.read() or b"").decode(errors="replace")[:200]
            if e.code in (401, 403):
                raise SystemExit(red("auth failed (%d) — check NSIN_API_KEY. %s" % (e.code, body)))
            if e.code == 404:
                raise SystemExit(red("%s: not found, or this key's user has no access." % domain))
            raise SystemExit(red("API error %d on %s: %s" % (e.code, domain, body)))
        except URLError as e:
            raise SystemExit(red("cannot reach %s: %s" % (API, e.reason)))

        batch = page.get("rows") or []
        if total is None:
            total = int(page.get("total") or 0)
            P.log(dim("  registry: %d cached entries" % total))
        for r in batch:
            k = (r.get("host", ""), r.get("path", ""), r.get("query", ""))
            if k in seen:
                continue
            seen.add(k)
            rows.append(r)
        offset += len(batch)
        if not batch or offset >= (total or 0):
            break
        P.set(["  %s %s" % (Progress.bar(offset / max(total, 1)),
                            dim("listing %d/%d" % (offset, total)))])
    return rows

# ---------- naming ----------


def dest_for(domain, row):
    """Map a cached URL onto a path under <OUTDIR>/<domain>/.

    Hostnames other than the domain itself (subdomains, wildcard records) get
    their own subfolder so two `/logo.png` from different hosts don't fight.
    """
    host = (row.get("host") or domain).lower()
    parts = [OUTDIR, domain]
    if host and host != domain:
        parts.append(re.sub(r'[<>:"\\|?*/]', "_", host))
    for p in unquote(row.get("path") or "").split("/"):
        p = p.strip().replace("\x00", "")
        if p in ("", ".", ".."):
            continue
        parts.append(re.sub(r'[<>:"\\|?*]', "_", p)[:180])
    if len(parts) == 2 or (row.get("path") or "/").endswith("/"):
        parts.append("index.html")
    return os.path.join(*parts)


def url_for(row, domain):
    host = row.get("host") or domain
    url = "https://%s%s" % (host, row.get("path") or "/")
    q = row.get("query") or ""
    return url + "?" + q if q else url


def wants(row, domain):
    """Which rows to fetch. Everything, unless --images narrows it.

    With --images the registry gives us no content-type, so the extension is
    all we have up front. Extensionless paths (/media/1234, image proxies) are
    fetched anyway and judged on their response content-type — it's the only
    way those images get picked up at all.
    """
    if not IMAGES:
        return True
    path = row.get("path") or "/"
    if IMG_EXT.search(path):
        return True
    base = path.rsplit("/", 1)[-1]
    return bool(base) and "." not in base


def plan_dests(domain, rows):
    """Assign every row its file path up front, resolving collisions here.

    Now that we take everything, one path routinely appears many times with
    different query strings (/search?q=a, /search?q=b) — and they are genuinely
    different cached objects. Colliding names get a short digest of their
    query+variant appended, which is stable across runs; doing this in the main
    thread (rather than racing in the workers) keeps the mapping deterministic.
    """
    groups = {}
    for r in rows:
        groups.setdefault(dest_for(domain, r), []).append(r)
    plan = []
    for base, rs in groups.items():
        if len(rs) == 1:
            plan.append((rs[0], base))
            continue
        root, ext = os.path.splitext(base)
        for r in rs:
            ident = "%s|%s" % (r.get("query") or "", r.get("variant") or "")
            tag = hashlib.sha1(ident.encode()).hexdigest()[:8]
            plan.append((r, "%s~%s%s" % (root, tag, ext)))
    return plan

# ---------- download ----------


class Stats:
    def __init__(self, total):
        self.total = total
        self.done = self.saved = self.skipped = self.failed = 0
        self.bytes = 0
        self.hits = 0
        self.start = time.monotonic()
        self.lock = threading.Lock()


def fetch(row, dest, domain, st, taken):
    url = url_for(row, domain)
    err, body, ctype, cache = "unknown", b"", "", ""
    for attempt in range(3):
        try:
            req = Request(url, headers={
                "Accept": "*/*",
                # identity: we want the stored bytes, not a gzip frame to unwrap.
                "Accept-Encoding": "identity",
                "User-Agent": "Mozilla/5.0 (nsin-dump-images)",
            })
            with urlopen(req, timeout=45, context=CTX) as r:
                body = r.read()
                ctype = (r.headers.get("Content-Type") or "").split(";")[0].strip().lower()
                cache = (r.headers.get("X-Cache") or r.headers.get("Cf-Cache-Status") or "").upper()
            err = None
            break
        except HTTPError as e:
            err = "HTTP %d" % e.code
            if e.code < 500:  # 404/403/410 won't get better on retry
                break
        except Exception as e:  # timeouts, TLS, DNS, resets
            err = type(e).__name__ + ": " + str(e)[:60]
        if attempt < 2:
            time.sleep(0.4 * (attempt + 1))

    if err:
        with st.lock:
            st.done += 1
            st.failed += 1
        P.log("  %s %s  %s" % (red("✗"), dim(row.get("path") or url), red(err)))
        return

    if not body or (IMAGES and not (ctype.startswith("image/") or IMG_EXT.search(dest))):
        with st.lock:
            st.done += 1
            st.skipped += 1
        return

    # Give extensionless URLs the extension their content-type implies.
    if not os.path.splitext(dest)[1] and ctype in EXT_FOR_CTYPE:
        dest += EXT_FOR_CTYPE[ctype]

    with st.lock:  # two variants of one URL can still collide on disk
        final, n = dest, 1
        while final in taken:
            root, ext = os.path.splitext(dest)
            final = "%s (%d)%s" % (root, n, ext)
            n += 1
        taken.add(final)

    os.makedirs(os.path.dirname(final) or ".", exist_ok=True)
    with open(final, "wb") as fh:
        fh.write(body)

    with st.lock:
        st.done += 1
        st.saved += 1
        st.bytes += len(body)
        if cache.startswith("HIT"):
            st.hits += 1
    P.log("  %s %s  %s" % (green("✓"),
                           os.path.relpath(final, OUTDIR),
                           dim("%s%s" % (human(len(body)), " · " + cache.lower() if cache else ""))))


def render(domain, st, idx, ndom):
    frac = st.done / st.total if st.total else 1.0
    el = time.monotonic() - st.start
    rate = st.bytes / el if el > 0.5 else 0
    eta = (st.total - st.done) / (st.done / el) if st.done and el > 0.5 else None
    head = "  %s %s  %s" % (Progress.bar(frac), bold(domain),
                            dim("[%d/%d]" % (idx, ndom)) if ndom > 1 else "")
    tail = "     %s  %s  %s  %s" % (
        "%3d%% %d/%d" % (frac * 100, st.done, st.total),
        green("↓%d" % st.saved) + ("  " + yellow("⊘%d" % st.skipped) if st.skipped else "")
        + ("  " + red("✗%d" % st.failed) if st.failed else ""),
        dim("%s · %s/s" % (human(st.bytes), human(rate))),
        dim("ETA " + clock(eta)),
    )
    P.set([head, tail])


def run(domain, idx, ndom):
    P.log("\n%s %s" % (cyan("▶"), bold(domain)))
    rows = list_keys(domain)
    plan = plan_dests(domain, [r for r in rows if wants(r, domain)])
    P.log(dim("  %d unique URLs, %d %s to fetch"
              % (len(rows), len(plan), "images" if IMAGES else "objects")))
    if not plan:
        return None

    st, taken = Stats(len(plan)), set()
    ticking = threading.Event()

    def ticker():  # keeps ETA/rate moving even while every worker is blocked
        while not ticking.wait(0.5):
            render(domain, st, idx, ndom)

    t = threading.Thread(target=ticker, daemon=True)
    t.start()
    try:
        with ThreadPoolExecutor(max_workers=WORKERS) as pool:
            for r, dest in plan:
                pool.submit(fetch, r, dest, domain, st, taken)
    finally:
        ticking.set()
        t.join(timeout=1)

    render(domain, st, idx, ndom)
    P.set([], force=True)
    P.log("  %s %s  %s" % (
        green("done"), bold(domain),
        dim("%d saved · %d skipped · %d failed · %s · %d edge hits"
            % (st.saved, st.skipped, st.failed, human(st.bytes), st.hits))))
    return st


try:
    results = []
    for i, d in enumerate(DOMAINS, 1):
        results.append((d, run(d, i, len(DOMAINS))))

    print()
    width = max(len(d) for d in DOMAINS)
    tot_saved = tot_bytes = tot_failed = 0
    for d, st in results:
        if st is None:
            print("  %-*s  %s" % (width, d, dim("nothing cached")))
            continue
        tot_saved += st.saved
        tot_bytes += st.bytes
        tot_failed += st.failed
        print("  %-*s  %s saved  %s%s" % (
            width, d, green(str(st.saved)), human(st.bytes),
            red("  %d failed" % st.failed) if st.failed else ""))
    if len(results) > 1:
        print("\n  %s %d files, %s%s" % (bold("total:"), tot_saved, human(tot_bytes),
                                         red(", %d failed" % tot_failed) if tot_failed else ""))
    print("  %s %s" % (dim("→"), OUTDIR))
except KeyboardInterrupt:
    P.close()
    sys.exit(130)
finally:
    P.close()
PY