"""Run the censor over the corpus and measure it. Given; do not change it.

    python3 experiments.py              all four parts
    python3 experiments.py corpus       A: what is in the traffic
    python3 experiments.py filters      B: the four strategies, measured
    python3 experiments.py evasion      C: which strategies still catch each evasion
    python3 experiments.py escalate     D: the censor's next move, and its cost

Paste what it prints into your report file. It also writes results.json.

This is the measurement harness. It is the only code that reads
corpus.truth(): a record is "targeted" when the site it really visits is on
corpus.BLOCKLIST, and "innocent" otherwise. For each strategy it counts

    targeted blocked    targeted records the filter blocked (true positives)
    targeted missed     targeted records that got through (false negatives)
    innocent blocked    innocent records blocked (false positives, the collateral)
    innocent sites hit  how many different innocent sites lost at least one visit
    bytes read          on average, how much of each record the filter must read
                        to decide: 4 for an address, the name for a lookup, the
                        whole ClientHello for SNI, the whole payload for keywords
    us per record       measured time per record, in microseconds, on this computer
"""

import argparse
import json
import random
import time
from collections import Counter

import corpus
import dpi
import evade


def the_censor(extra_hosts=(), extra_ips=()):
    ips = {corpus.resolve(h) for h in corpus.BLOCKLIST} | set(extra_ips)
    return dpi.Censor(hosts=list(corpus.BLOCKLIST) + list(extra_hosts), ips=ips, keywords=list(corpus.KEYWORDS))


def bytes_read(strategy, r):
    if strategy == "IP":
        return 4
    if strategy == "DNS":
        return len(r["dns"] or "")
    if strategy == "SNI":
        return r["hello_bytes"]
    return len(r["payload"]) // 2 if r["port"] == 443 else len(r["payload"])


def score(fn, records, truth, censor):
    targeted = [s in corpus.BLOCKLIST for s in truth]
    blocked = [bool(fn(r, censor)) for r in records]
    rows = range(len(targeted))
    tp = sum(1 for i in rows if blocked[i] and targeted[i])
    fn_ = sum(1 for i in rows if targeted[i] and not blocked[i])
    fp = sum(1 for i in rows if blocked[i] and not targeted[i])
    hit = Counter(truth[i] for i in rows if blocked[i] and not targeted[i])
    return {"targeted_blocked": tp, "targeted_missed": fn_, "innocent_blocked": fp,
            "innocent_sites_hit": len(hit), "collateral": dict(hit.most_common())}


def per_record_us(fn, records, censor):
    runs, total = 0, 0.0
    while total < 0.1 and runs < 200:
        start = time.perf_counter()
        for r in records:
            fn(r, censor)
        total += time.perf_counter() - start
        runs += 1
    return total / runs / len(records) * 1e6


def part_corpus(records, truth):
    print("## A. The corpus\n")
    sites_on = {}
    for i in range(len(truth)):
        r, s = records[i], truth[i]
        sites_on.setdefault(r["ip"], set()).add(s)
    shared = {ip: len(s) for ip, s in sites_on.items() if len(s) > 1}
    targeted = sum(s in corpus.BLOCKLIST for s in truth)
    stats = {"records": len(records), "sites": len(set(truth)), "addresses": len(sites_on),
             "shared_addresses": shared, "plain_http": sum(r["port"] == 80 for r in records),
             "no_lookup_seen": sum(r["dns"] is None for r in records),
             "already_ech": sum(r["sni"] in corpus.ECH_PUBLIC_NAME.values() for r in records),
             "targeted": targeted}
    print(f"{stats['records']} records, {stats['sites']} sites, {stats['addresses']} addresses.")
    for ip, n in sorted(shared.items()):
        print(f"  {ip} is shared by {n} sites")
    print(f"{stats['plain_http']} records are plain HTTP (port 80); the rest are HTTPS.")
    print(f"{stats['no_lookup_seen']} records have no DNS lookup seen (the answer was cached).")
    print(f"{stats['already_ech']} records already use ECH (SNI is a CDN's public name).")
    print(f"{targeted} records visit a blocked site ({', '.join(corpus.BLOCKLIST)}); "
          f"keywords: {', '.join(corpus.KEYWORDS)}.\n")
    return stats


def part_filters(records, truth):
    print("## B. Four strategies, measured\n")
    censor = the_censor()
    print(f"Blocked addresses (from public DNS): {', '.join(sorted(censor.ips))}\n")
    out = {}
    for name, fn in dpi.FILTERS.items():
        try:
            s = score(fn, records, truth, censor)
        except NotImplementedError:
            print(f" ({name} filter not written yet; left out of this table)")
            continue
        s["bytes_per_record"] = round(sum(bytes_read(name, r) for r in records) / len(records), 1)
        s["us_per_record"] = round(per_record_us(fn, records, censor), 3)
        out[name] = s
    print("| strategy | targeted blocked | targeted missed | innocent blocked | innocent sites hit | bytes read per record | us per record |")
    print("| :--- | --: | --: | --: | --: | --: | --: |")
    for name, s in out.items():
        print(f"| {name} | {s['targeted_blocked']} | {s['targeted_missed']} | {s['innocent_blocked']} | "
              f"{s['innocent_sites_hit']} | {s['bytes_per_record']} | {s['us_per_record']:.2f} |")
    print("\nInnocent sites each strategy blocked (visits blocked):")
    for name, s in out.items():
        hit = ", ".join(f"{site} ({n})" for site, n in s["collateral"].items()) or "none"
        print(f"  {name}: {hit}")
    print()
    return out


def evaded_world(records, truth, evasion, innocent_share=0.0, seed=10):
    """Every targeted record uses the evasion; optionally a share of innocent records does too."""
    rng = random.Random(seed)
    out = []
    for i in range(len(truth)):
        r, s = records[i], truth[i]
        uses = s in corpus.BLOCKLIST or rng.random() < innocent_share
        out.append(evasion(dict(r)) if uses else dict(r))
    return out


def written_evasions():
    ok = {}
    probe = corpus.records()[0]
    for name, fn in evade.EVASIONS.items():
        try:
            fn(probe)
            ok[name] = fn
        except NotImplementedError:
            print(f"  ({name} not written yet; skipped)")
    if "hide_sni" in ok and "use_doh" in ok:
        ok["hide_sni + use_doh"] = lambda r: evade.use_doh(evade.hide_sni(r))
    return ok


def usable_filters(records, censor):
    """The strategies that are actually written, so an extra-credit filter left
    unwritten does not drop the whole part."""
    ok = {}
    for name, fn in dpi.FILTERS.items():
        try:
            fn(records[0], censor)
        except NotImplementedError:
            print(f" ({name} filter not written yet; left out of this part)")
            continue
        ok[name] = fn
    return ok

def part_evasion(records, truth):
    print("## C. Evasion: share of targeted visits each strategy still blocks\n")
    evasions = written_evasions()
    if not evasions:
        raise NotImplementedError("hide_sni")
    censor = the_censor()
    filters = usable_filters(records, censor)
    names = list(filters)
    t_idx = [i for i, s in enumerate(truth) if s in corpus.BLOCKLIST]
    out, lines = {}, []
    for label, fn in [("none", dict)] + list(evasions.items()):
        world = evaded_world(records, truth, fn)
        caught = {name: sum(bool(f(world[i], censor)) for i in t_idx) for name, f in filters.items()}
        anyf = sum(any(f(world[i], censor) for f in filters.values()) for i in t_idx)
        out[label] = {**{k: v / len(t_idx) for k, v in caught.items()}, "any": anyf / len(t_idx)}
        lines.append((label, caught, anyf))
    print("| evasion | " + " | ".join(names) + " | any strategy written |")
    print("| :--- | " + " | ".join("--:" for _ in names) + " | --: |")
    for label, caught, anyf in lines:
        print(f"| {label} | " + " | ".join(f"{caught[n] / len(t_idx):.0%}" for n in names)
              + f" | {anyf / len(t_idx):.0%} |")
    print(f"\n({len(t_idx)} targeted visits.)\n")
    return out


def part_escalate(records, truth):
    print("## D. Escalation: the censor's next move, and what it costs\n")
    evasions = written_evasions()
    base = the_censor()
    rows, out = [], {}
    # Each row: the situation, the censor's response (filter and settings), and what it used before.
    filters = usable_filters(records, base)
    if "hide_sni" in evasions and "SNI" in filters:
        world = evaded_world(records, truth, evasions["hide_sni"])
        rows.append(("ECH wherever the provider offers it; censor adds the ECH public name to its SNI list",
                     filters["SNI"], the_censor(extra_hosts=corpus.ECH_PUBLIC_NAME.values()), world, filters["SNI"]))
    if "use_doh" in evasions:
        world = evaded_world(records, truth, evasions["use_doh"])
        rows.append(("DNS over HTTPS; censor switches from DNS to address blocking",
                     dpi.ip_filter, base, world, dpi.dns_filter))
    if "via_proxy" in evasions:
        world = evaded_world(records, truth, evasions["via_proxy"], innocent_share=0.1)
        rows.append(("a proxy, also used by 10% of innocent visits; censor adds the proxy's address",
                     dpi.ip_filter, the_censor(extra_ips=[corpus.PROXY["ip"]]), world, dpi.ip_filter))
    if not rows:
        raise NotImplementedError("hide_sni")
    results = []
    for label, fn, censor, world, before in rows:
        s = score(fn, world, truth, censor)
        s["innocent_added"] = s["innocent_blocked"] - score(before, world, truth, base)["innocent_blocked"]
        out[label] = s
        results.append((label, s))
    print("| situation and response | targeted blocked | targeted missed | innocent blocked | added by this response | innocent sites hit |")
    print("| :--- | --: | --: | --: | --: | --: |")
    for label, s in results:
        print(f"| {label} | {s['targeted_blocked']} | {s['targeted_missed']} | {s['innocent_blocked']} | "
              f"{s['innocent_added']} | {s['innocent_sites_hit']} |")
    print()
    return out


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("part", nargs="?", choices=["corpus", "filters", "evasion", "escalate"], help="run one part (default: all)")
    args = ap.parse_args()
    records, truth = corpus.records(), corpus.truth()
    try:
        with open("results.json") as f:
            out = json.load(f)
    except (OSError, ValueError):
        out = {}
    parts = [("corpus", part_corpus), ("filters", part_filters), ("evasion", part_evasion), ("escalate", part_escalate)]
    for name, run in parts:
        if args.part not in (None, name):
            continue
        try:
            out[name] = run(records, truth)
        except NotImplementedError as e:
            print(f"\n## {name}: skipped, {e} is not written yet (python3 selfcheck.py shows what is left)\n")
    try:
        with open("results.json", "w") as f:
            json.dump(out, f, indent=1)
    except OSError as exc:
        print(f"could not write results.json: {exc}")
        return 1
    print("wrote results.json")


if __name__ == "__main__":
    main()