"""Second round of the September probe, and the comparison. Run: python -B run.py [--compare-only]
Raw rows: ~/data/llms-two-months/round2.jsonl. Writes results.json next to this file."""
import collections
import hashlib
import importlib.util
import json
import os
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed

import rule

HERE = os.path.dirname(os.path.abspath(__file__))
PROBE = os.path.join(HERE, "..", "llms-txt", "probe.py")
DATA = os.path.expanduser("~/data/llms-two-months")
SEPT = os.path.expanduser("~/data/llms-txt/results.jsonl")
OUT = os.path.join(DATA, "round2.jsonl")

sys.argv = sys.argv[:1] + [a for a in sys.argv[1:] if a == "--compare-only"]
COMPARE_ONLY = "--compare-only" in sys.argv
sys.argv = sys.argv[:1]
spec = importlib.util.spec_from_file_location("probe", PROBE)
probe = importlib.util.module_from_spec(spec)
spec.loader.exec_module(probe)


def load(path):
    return {r["domain"]: r for r in map(json.loads, open(path))}


def fetch(sept):
    os.makedirs(DATA, exist_ok=True)
    done = load(OUT) if os.path.exists(OUT) else {}
    todo = [(r["rank"], d) for d, r in sept.items() if d not in done]
    with open(OUT, "a") as fo, ThreadPoolExecutor(max_workers=32) as ex:
        for fut in as_completed([ex.submit(probe.probe, r, d) for r, d in todo]):
            fo.write(json.dumps(fut.result()) + "\n")
            fo.flush()


def pair(sept, now, test, keys):
    g = [d for d in keys if not test(sept[d]) and test(now[d])]
    l = [d for d in keys if test(sept[d]) and not test(now[d])]
    net, lo, hi = rule.paired_interval(len(g), len(l), len(keys))
    return {"base": sum(test(sept[d]) for d in keys), "now": sum(test(now[d]) for d in keys),
            "gained": len(g), "lost": len(l), "net": net, "lo": lo, "hi": hi, "gained_domains": sorted(g),
            "lost_domains": sorted(l)}


def main():
    sept = load(SEPT)
    if not COMPARE_ONLY:
        fetch(sept)
    now = load(OUT)
    sha = hashlib.sha256(open(PROBE, "rb").read()).hexdigest()
    keys = sorted(d for d in sept if d in now and rule.answered(sept[d]) and rule.answered(now[d]))
    llms = lambda r: bool(r["llms_ok"])
    p = pair(sept, now, llms, keys)
    retained = (p["base"] - p["lost"]) / p["base"] if p["base"] else 0
    rs = sum(sept[d]["robots_status"] == 200 for d in keys)
    rn = sum(now[d]["robots_status"] == 200 for d in keys)
    voids = {"V1": len(keys) < rule.MIN_PAIRED, "V2": retained < rule.MIN_RETAINED,
             "V3": abs(rn - rs) / rs > rule.MAX_ROBOTS_MOVE if rs else True, "V4": sha != rule.PROBE_SHA256}
    v = "void" if any(voids.values()) else rule.verdict(p["lo"], p["hi"], p["base"])
    firsts = collections.Counter(now[d]["llms_first"] for d in p["gained_domains"])
    res = {"paired": len(keys), "answered_sept": sum(rule.answered(r) for r in sept.values()),
           "answered_now": sum(rule.answered(r) for r in now.values()), "llms": p, "bar": rule.BAR * p["base"],
           "retained": retained, "robots_200": [rs, rn], "voids": voids, "verdict": v,
           "printed_only": {
               "all_1000": {"llms_sept": sum(map(llms, sept.values())), "llms_now": sum(map(llms, now.values())),
                            "blocks_sept": sum(rule.blocks_any(r, probe.CRAWLERS) for r in sept.values()),
                            "blocks_now": sum(rule.blocks_any(r, probe.CRAWLERS) for r in now.values())},
               "blocks_any": pair(sept, now, lambda r: rule.blocks_any(r, probe.CRAWLERS), keys),
               "per_crawler": {c: {k: x for k, x in pair(sept, now, lambda r, c=c: rule.blocks_any(r, [c]), keys).items()
                                   if not k.endswith("domains")} for c in probe.CRAWLERS},
               "top_100": pair(sept, now, llms, [d for d in keys if sept[d]["rank"] <= 100]),
               "gained_first_lines": firsts.most_common(40),
               "gained_lengths": sorted(now[d]["llms_len"] for d in p["gained_domains"])}}
    json.dump(res, open(os.path.join(HERE, "results.json"), "w"), indent=1)
    print(v, len(keys), p["base"], p["gained"], p["lost"], p["net"], round(p["lo"], 1), round(p["hi"], 1), voids)


main()