"""Rule steps 1-5, printed-only captures. Usage: python -B probe.py SCAN_JSONL BODIES_DIR OUT_JSON"""
import json, os, sys, time
from concurrent.futures import ThreadPoolExecutor
import requests, urllib3
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import rule

urllib3.disable_warnings()
UA = "untilnextsession-research/1.0 (+https://untilnextsession.com/experiments/; one MCP initialize, nothing further)"
RANK = {"SPEAKS": 2, "LOCKED": 1, "NOTHING": 0}


def ask(url):
    """(class, status, server name). One POST, redirects not followed past what requests does for POST."""
    try:
        with requests.post(url, json=rule.INITIALIZE, timeout=rule.TIMEOUT_S, verify=False, stream=True,
                           headers={"User-Agent": UA, "Accept": "application/json, text/event-stream"}) as r:
            body, t0 = b"", time.time()
            for chunk in r.iter_content(16384):
                body += chunk
                # An event stream may stay open: stop at the first full message or the clock.
                if len(body) >= rule.MAX_BODY or time.time() - t0 > rule.TIMEOUT_S or (
                        b"data:" in body and b"\n\n" in body):
                    break
            headers = {k.lower(): v for k, v in r.headers.items()}
            return rule.classify(r.status_code, headers, body), r.status_code, rule.server_name(body)
    except Exception as e:
        return "NOTHING", None, type(e).__name__


def first_capture(domain, final):
    """Printed only: first 200 capture of the card path, from the Internet Archive."""
    try:
        host = final.split("/")[2].split(":")[0]
        r = requests.get("https://web.archive.org/cdx/search/cdx", timeout=40, headers={"User-Agent": UA},
                         params={"url": f"{host}/.well-known/mcp/server-card.json", "limit": 1,
                                 "filter": "statuscode:200", "output": "json"})
        rows = r.json() if r.status_code == 200 and r.text.strip() else []
        return rows[1][1][:8] if len(rows) > 1 else None
    except Exception as e:
        return "err:" + type(e).__name__


def main(scan, bodies, out):
    rows = []
    for line in open(scan):
        r = json.loads(line)
        if r["reachable"] and r["paths"]["mcp_card"]["ok"]:
            p = r["paths"]["mcp_card"]
            card = json.loads(open(os.path.join(bodies, p["sha"] + ".body"), "rb").read().decode("utf-8-sig"))
            probe, tmpl = rule.urls_of(card)
            rows.append({"domain": r["domain"], "rank": r["rank"], "a": r["a"], "b": r["b"], "final": p["final"],
                         "urls": probe, "templated": tmpl, "answers": {}})
    urls = sorted({u for r in rows for u in r["urls"]})
    with ThreadPoolExecutor(8) as ex:
        first = dict(zip(urls, ex.map(ask, urls)))
    missed = [u for u, a in first.items() if a[0] == "NOTHING"]
    print(f"{len(urls)} urls, {len(missed)} NOTHING, retry in {rule.RETRY_MIN} min", file=sys.stderr, flush=True)
    time.sleep(rule.RETRY_MIN * 60)
    with ThreadPoolExecutor(8) as ex:
        second = dict(zip(missed, ex.map(ask, missed)))
    best = {u: max((a for a in (first[u], second.get(u)) if a), key=lambda a: RANK[a[0]]) for u in urls}
    for r in rows:
        r["answers"] = {u: best[u] for u in r["urls"]}
        r["door"] = any(best[u][0] != "NOTHING" for u in r["urls"])
        r["first_capture"] = first_capture(r["domain"], r["final"])
        time.sleep(1.5)
    a = [r for r in rows if r["a"]]
    tested = [r for r in a if r["urls"]]
    k = sum(r["door"] for r in tested)
    res = {"frame": len(a), "no_url": sum(1 for r in a if not r["urls"] and not r["templated"]),
           "only_templated": sum(1 for r in a if not r["urls"] and r["templated"]),
           "tested": len(tested), "door": k, "ci": rule.wilson(k, len(tested)),
           "verdict": "void" if len(tested) < 30 else rule.verdict(k, len(tested)), "rows": rows}
    json.dump(res, open(out, "w"), indent=1)
    print(res["verdict"], k, len(tested), file=sys.stderr)


if __name__ == "__main__":
    main(*sys.argv[1:4])