"""Re-fetch, through a rotating proxy, every domain the direct run could not settle.

A 403 or a timeout from one IP is not evidence about a site's robots.txt: it may
be the network refusing the request rather than the site refusing the agent.
Every such domain is fetched again from a fresh proxy IP with both user-agents,
once with a retry, and the proxy result replaces the direct one whenever it
produces a 200. Domains that answered 200 or a clean 404 to either agent are not
re-fetched.

The proxy URL is read from the environment (PROXY_URL) and is never written to
any output file. curl picks it up from https_proxy/http_proxy, so the credential
does not appear in a process argument list either.
"""
import gzip, hashlib, json, os, subprocess, sys, tempfile, time
from concurrent.futures import ThreadPoolExecutor

PROXY = os.environ.get("PROXY_URL")
if not PROXY:
    sys.exit("set PROXY_URL")
ENV = dict(os.environ, https_proxy=PROXY, http_proxy=PROXY, ALL_PROXY=PROXY)

UA_RESEARCH = ("FentnerResearchBot/1.0 (+https://fentner.com/research; "
               "robots.txt measurement study; contact: research@fentner.com)")
UA_BROWSER = ("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 "
              "(KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36")
MAXB = 300000
SOFT_RETRY = {403, 429, 500, 502, 503, 504, 520, 521, 522, 523, 524}
SOFT = {-1, 0, 401, 403, 406, 408, 409, 425, 429}


def needs_retry(d):
    r, b = d["r_status"], d["b_status"]
    if r == 200 or b == 200:
        # settled, unless one agent was refused in a way a fresh IP might change
        return r in SOFT or b in SOFT
    if r in (404, 410) and b in (404, 410):
        return False          # the site genuinely has no robots.txt
    return True


def get(url, ua, attempts=2):
    """Fetch once, and retry from a fresh proxy IP only when the first attempt
    got an answer that a different IP might change (a refusal or a rate limit).
    A host that does not answer at all is not retried: a second 30-second wait
    buys nothing and the run has thousands of dead domains in it."""
    last = {"status": -1, "url": url, "ctype": "", "body": ""}
    for attempt in range(attempts):
        if attempt and last["status"] not in SOFT_RETRY:
            break
        fd, path = tempfile.mkstemp(prefix="px_")
        os.close(fd)
        cmd = ["curl", "-s", "-L", "-k", "--compressed",
               "--connect-timeout", "8", "--max-time", "20", "--max-redirs", "5",
               "--max-filesize", str(MAXB * 4),
               "-A", ua, "-H", "Accept: */*", "-H", "Accept-Language: en-GB,en;q=0.9",
               "-o", path, "-w", "%{http_code}\t%{url_effective}\t%{content_type}", url]
        try:
            p = subprocess.run(cmd, capture_output=True, timeout=60, env=ENV)
            meta = p.stdout.decode("utf-8", "replace").split("\t")
            code = int(meta[0]) if meta and meta[0].isdigit() else -1
            with open(path, "rb") as f:
                body = f.read(MAXB).decode("utf-8", "replace")
            last = {"status": code, "url": meta[1] if len(meta) > 1 else url,
                    "ctype": meta[2][:120] if len(meta) > 2 else "", "body": body}
        except Exception:
            pass
        finally:
            try: os.unlink(path)
            except OSError: pass
        if last["status"] == 200:
            return last
    return last


def sha(s):
    return hashlib.sha256(s.encode("utf-8", "replace")).hexdigest()[:16]


def job(d):
    u = d.get("url") or f"https://{d['domain']}/robots.txt"
    r = get(u, UA_RESEARCH)
    if r["status"] in (-1, 0):
        u2 = f"https://www.{d['domain']}/robots.txt"
        r2 = get(u2, UA_RESEARCH)
        if r2["status"] not in (-1, 0):
            u, r = u2, r2
    b = get(u, UA_BROWSER)
    if r["status"] != 200 and b["status"] != 200:
        # the proxy did no better; keep the direct result, record the attempt
        out = dict(d)
        out["proxy_r_status"] = r["status"]
        out["proxy_b_status"] = b["status"]
        out["fetched_via"] = "direct"
        return out
    out = {"rank": d["rank"], "domain": d["domain"], "url": u,
           "r_status": r["status"], "r_url": r["url"], "r_ctype": r["ctype"],
           "b_status": b["status"], "b_url": b["url"], "b_ctype": b["ctype"],
           "r_sha": sha(r["body"]), "b_sha": sha(b["body"]),
           "r_len": len(r["body"]), "b_len": len(b["body"]),
           "r_body": r["body"],
           "direct_r_status": d["r_status"], "direct_b_status": d["b_status"],
           "fetched_via": "proxy"}
    out["b_body"] = "" if out["r_sha"] == out["b_sha"] else b["body"]
    return out


skip = set(json.load(open(os.environ["SKIP_FILE"]))) if os.environ.get("SKIP_FILE") else set()
targets = []
for p in sys.argv[2:]:
    for line in gzip.open(p, "rt"):
        d = json.loads(line)
        if needs_retry(d) and d["domain"] not in skip:
            targets.append(d)
print("to re-fetch via proxy:", len(targets), flush=True)
t0 = time.time()
with gzip.open(sys.argv[1], "wt") as f, ThreadPoolExecutor(max_workers=20) as ex:
    for i, r in enumerate(ex.map(job, targets)):
        f.write(json.dumps(r) + "\n")
        if i % 250 == 0:
            f.flush()
            print(i, len(targets), f"{time.time()-t0:.0f}s", flush=True)
print("done", time.time() - t0, flush=True)