#!/usr/bin/env python3
"""Cache conformance for HTTP delivery (extensions/http-delivery.md §9).

Runs tests.json against an HTTP cache that fronts an origin serving the objects of
wai/http-conformance/, and reports, per test, whether the cache met it. Standard library
only, and ../claim-conformance/_blake3.py to check each object's bytes against its
content hash; it shares no code with wai-rs.

    python3 driver.py stage <http-conformance dir> <root>
        lay out the files an origin serves for the tests under <root>
    python3 driver.py run --base URL --root <root> [--origin URL] [--label L]
                          [--version V] [--config FILE] [--report report.json]
        run the tests against the cache at URL (the origin itself, for a control)

`run` exits 1 when a MUST test fails, and prints `FAIL <id>: <detail>` for each. With
--origin, the origin's hit counter (`wai http origin --count-hits`) lets the informative
tests see whether a request reached the origin. A report names the cache by its label
and version only.
"""
import hashlib, http.client, json, os, shutil, sys, time, urllib.parse

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "claim-conformance"))
from _blake3 import blake3  # noqa: E402

DOMAIN = b"moq-jwp:object\x01"


def cid_of(b):
    return blake3(DOMAIN + b).hex()


# ------------------------------------------------------------------- staging

def stage(corpus, root):
    """The origin's files: the corpus objects and head proofs, an alias, a live-edge
    part, a receipt index, a receipt file, a playlist receipt, a key log and a
    parameter set named by its digest."""
    if os.path.exists(root):
        shutil.rmtree(root)
    os.makedirs(os.path.join(root, "objects"))
    listing = json.load(open(os.path.join(corpus, "objects.json")))
    for oid, o in listing.items():
        for suffix in ("", ".hp"):
            shutil.copyfile(os.path.join(corpus, o["path"] + suffix), os.path.join(root, o["path"] + suffix))
    def put(rel, data):
        p = os.path.join(root, rel)
        os.makedirs(os.path.dirname(p), exist_ok=True)
        open(p, "wb").write(data)
    o8 = open(os.path.join(corpus, listing["O8"]["path"]), "rb").read()
    o9b = open(os.path.join(corpus, listing["O9b"]["path"]), "rb").read()
    put("alias/segment-8.wai", o8)
    put("live/part-1.wai", o9b)
    index = json.dumps({"kind": "wai-receipt-index", "object": listing["O3"]["cid"], "delivery": [], "claims": []}, sort_keys=True).encode()
    put(listing["O3"]["path"] + ".receipts", index)
    fixture = json.dumps({"kind": "cache-conformance-fixture", "of": "a receipt file"}, sort_keys=True).encode()
    put("objects/rcpt/%s.json" % blake3(fixture).hex(), fixture)
    playlist = json.dumps({"kind": "cache-conformance-fixture", "of": "a playlist receipt"}, sort_keys=True).encode()
    put("objects/mr/%s.json" % blake3(playlist).hex(), playlist)
    put(".well-known/wai-keys", json.dumps({"kind": "cache-conformance-fixture", "of": "a key log"}).encode())
    pset = open(os.path.join(corpus, "..", "wai-rs", "tests", "vectors", "int_mlicpp", "model.wmb"), "rb").read()
    put("sets/" + hashlib.sha256(pset).hexdigest(), pset)
    manifest = {
        "objects": {oid: {"path": o["path"], "cid": o["cid"], "n": o["n"]} for oid, o in listing.items()},
        "alias": {"path": "alias/segment-8.wai", "of": "O8"},
        "live": {"path": "live/part-1.wai", "of": "O9b"},
        "index": listing["O3"]["path"] + ".receipts",
        "receipt": "objects/rcpt/%s.json" % blake3(fixture).hex(),
        "playlist_receipt": "objects/mr/%s.json" % blake3(playlist).hex(),
        "key_log": ".well-known/wai-keys",
        "set": "sets/" + hashlib.sha256(pset).hexdigest(),
    }
    open(os.path.join(root, "stage.json"), "w").write(json.dumps(manifest, indent=1) + "\n")
    print("staged %d objects and their sidecars under %s" % (len(listing), root))


# ------------------------------------------------------------------- requests

class Response:
    def __init__(self, status, headers, body):
        self.status, self.body = status, body
        self.headers = {}
        for k, v in headers:
            self.headers.setdefault(k.lower(), []).append(v)

    def h(self, name):
        v = self.headers.get(name.lower())
        return None if v is None else ", ".join(v)


def request(base, path, headers=None, method="GET"):
    u = urllib.parse.urlsplit(base)
    conn = http.client.HTTPConnection(u.hostname, u.port or 80, timeout=30)
    prefix = u.path.rstrip("/")
    h = {"Connection": "close"}
    h.update(headers or {})
    conn.request(method, prefix + "/" + path.lstrip("/"), headers=h)
    r = conn.getresponse()
    body = r.read()
    out = Response(r.status, r.getheaders(), body)
    conn.close()
    return out


def hits(origin, path):
    if origin is None:
        return None
    r = request(origin, "/__hits?path=" + urllib.parse.quote("/" + path.lstrip("/")))
    return int(r.body) if r.status == 200 else None


# --------------------------------------------------------------------- tests

class Fail(Exception):
    pass


def need(cond, why):
    if not cond:
        raise Fail(why)


def run_tests(base, root, origin):
    st = json.load(open(os.path.join(root, "stage.json")))
    obj = {k: open(os.path.join(root, v["path"]), "rb").read() for k, v in st["objects"].items()}
    for k, v in st["objects"].items():
        if cid_of(obj[k]) != v["cid"]:
            raise SystemExit("the staged object %s does not hash to its cid" % k)
    path = lambda k: st["objects"][k]["path"]
    etag = lambda k: '"%s"' % st["objects"][k]["cid"]
    results = {}

    def whole_ok(r, k, what):
        need(r.status == 200, "%s: status %d" % (what, r.status))
        need(r.h("content-encoding") in (None, "identity"), "%s: Content-Encoding %s" % (what, r.h("content-encoding")))
        need(r.body == obj[k], "%s: %d bytes that are not the object's %d" % (what, len(r.body), len(obj[k])))

    def range_ok(r, k, a, b, what, allow_200=False):
        n = len(obj[k])
        if allow_200 and r.status == 200:
            whole_ok(r, k, what)
            return
        need(r.status == 206, "%s: status %d" % (what, r.status))
        need(r.h("content-encoding") in (None, "identity"), "%s: Content-Encoding %s" % (what, r.h("content-encoding")))
        need(r.h("content-range") == "bytes %d-%d/%d" % (a, b, n), "%s: Content-Range %s" % (what, r.h("content-range")))
        need(r.body == obj[k][a:b + 1], "%s: the bytes are not the object's [%d, %d]" % (what, a, b))

    def c1():
        whole_ok(request(base, path("O3")), "O3", "cold")
        whole_ok(request(base, path("O3")), "O3", "again")

    def c2():
        r = request(base, path("O2"))
        whole_ok(r, "O2", "GET")
        need(r.h("etag") == etag("O2"), "ETag %s" % r.h("etag"))
        need((r.h("content-type") or "").startswith("application/wai"), "Content-Type %s" % r.h("content-type"))
        need(r.h("cache-control") == "public, max-age=31536000, immutable, no-transform", "Cache-Control %s" % r.h("cache-control"))
        need(r.h("content-length") == str(len(obj["O2"])), "Content-Length %s" % r.h("content-length"))

    def c3():
        request(base, path("O4"))
        range_ok(request(base, path("O4"), {"Range": "bytes=1009-3000"}), "O4", 1009, 3000, "range")

    def c4():
        n = len(obj["O7"])
        range_ok(request(base, path("O7"), {"Range": "bytes=1009-2999"}), "O7", 1009, 2999, "cold range", allow_200=True)
        range_ok(request(base, path("O7"), {"Range": "bytes=0-1008"}), "O7", 0, 1008, "head range", allow_200=True)
        need(n > 0, "empty")

    def c5():
        n = len(obj["O9c"])
        request(base, path("O9c"))
        range_ok(request(base, path("O9c"), {"Range": "bytes=-100"}), "O9c", n - 100, n - 1, "suffix")
        range_ok(request(base, path("O9c"), {"Range": "bytes=%d-" % (n - 50)}), "O9c", n - 50, n - 1, "open")

    def c6():
        n = len(obj["O9c"])
        r = request(base, path("O9c"), {"Range": "bytes=%d-" % n})
        need(r.status == 416, "status %d" % r.status)
        need(r.h("content-range") == "bytes */%d" % n, "Content-Range %s" % r.h("content-range"))

    def c7():
        for enc in ("gzip, deflate, br", "br", "zstd, gzip"):
            whole_ok(request(base, path("O9a"), {"Accept-Encoding": enc}), "O9a", "Accept-Encoding %s" % enc)

    def c8():
        request(base, path("OD"))
        range_ok(request(base, path("OD"), {"Range": "bytes=1009-3593", "If-Range": etag("OD")}), "OD", 1009, 3593, "matching If-Range")
        whole_ok(request(base, path("OD"), {"Range": "bytes=1009-3593", "If-Range": "W/" + etag("OD")}), "OD", "weak If-Range")
        whole_ok(request(base, path("OD"), {"Range": "bytes=1009-3593", "If-Range": '"%s"' % ("0" * 64)}), "OD", "other If-Range")

    def c9():
        a = request(base, path("O5"), {"User-Agent": "conformance-a", "Accept-Language": "en", "Accept": "*/*"})
        b = request(base, path("O5"), {"User-Agent": "conformance-b", "Accept-Language": "fr", "Accept": "application/wai"})
        whole_ok(a, "O5", "first")
        whole_ok(b, "O5", "second")
        need(a.h("vary") is None and b.h("vary") is None, "Vary %s / %s" % (a.h("vary"), b.h("vary")))

    def c10():
        p = path("O3") + ".hp"
        r = request(base, p)
        need(r.status == 200 and r.body == open(os.path.join(root, p), "rb").read(), "the head proof differs (status %d)" % r.status)
        need(r.h("content-encoding") in (None, "identity"), "Content-Encoding %s" % r.h("content-encoding"))
        need(r.h("cache-control") == "public, max-age=31536000, immutable, no-transform", "Cache-Control %s" % r.h("cache-control"))

    def c11():
        r = request(base, st["alias"]["path"])
        whole_ok(r, st["alias"]["of"], "alias")
        need(r.h("etag") == etag(st["alias"]["of"]), "ETag %s" % r.h("etag"))

    def c12():
        k, p = st["live"]["of"], st["live"]["path"]
        r = request(base, p)
        whole_ok(r, k, "live part")
        cc = r.h("cache-control") or ""
        need("max-age=31536000" not in cc and "max-age=" in cc, "Cache-Control %s" % cc)
        range_ok(request(base, p, {"Range": "bytes=10-99"}), k, 10, 99, "live range", allow_200=True)

    def c13():
        r = request(base, path("O1"), {"Origin": "https://reader.example"})
        whole_ok(r, "O1", "with Origin")
        need(r.h("access-control-allow-origin") == "*", "Access-Control-Allow-Origin %s" % r.h("access-control-allow-origin"))
        need("origin" not in (r.h("vary") or "").lower(), "Vary %s" % r.h("vary"))

    def r1():
        p = st["index"]
        first = request(base, p)
        need(first.status == 200, "status %d" % first.status)
        changed = json.dumps({"kind": "wai-receipt-index", "changed_at": time.time()}).encode()
        full = os.path.join(root, p)
        old = open(full, "rb").read()
        open(full, "wb").write(changed)
        try:
            second = request(base, p)
        finally:
            open(full, "wb").write(old)
        need(second.status == 200 and second.body == changed, "the index came back unchanged")

    def fetch_kept(p, cc=None):
        r = request(base, p)
        need(r.status == 200 and r.body == open(os.path.join(root, p), "rb").read(), "%s differs (status %d)" % (p, r.status))
        need(r.h("content-encoding") in (None, "identity"), "Content-Encoding %s" % r.h("content-encoding"))
        if cc is not None:
            need(r.h("cache-control") == cc, "Cache-Control %s" % r.h("cache-control"))
        return r

    def kept(p, cc=None):
        fetch_kept(p, cc)

    imm = "public, max-age=31536000, immutable, no-transform"

    def r5():
        r = fetch_kept(st["set"], imm)
        need(hashlib.sha256(r.body).hexdigest() == st["set"].rsplit("/", 1)[1], "the set does not hash to its name")

    def s1():
        if origin is None:
            return "not run: no hit counter"
        p = path("O8")
        request(base, p)
        before = hits(origin, p)
        request(base, p)
        return "stored" if hits(origin, p) == before else "not stored"

    def s2():
        if origin is None:
            return "not run: no hit counter"
        p = path("O9b")
        request(base, p)
        before = hits(origin, p)
        request(base, p, {"Range": "bytes=100-199"})
        return "answered from store" if hits(origin, p) == before else "forwarded"

    table = {"C1": c1, "C2": c2, "C3": c3, "C4": c4, "C5": c5, "C6": c6, "C7": c7, "C8": c8, "C9": c9,
             "C10": c10, "C11": c11, "C12": c12, "C13": c13, "R1": r1,
             "R2": lambda: kept(st["receipt"], imm), "R3": lambda: kept(st["playlist_receipt"], imm),
             "R4": lambda: kept(st["key_log"], "max-age=300, no-transform"), "R5": r5, "S1": s1, "S2": s2}
    tests = json.load(open(os.path.join(HERE, "tests.json")))["tests"]
    out = []
    for t in tests:
        try:
            detail = table[t["id"]]()
            out.append({"id": t["id"], "name": t["name"], "level": t["level"], "pass": True, "detail": detail})
        except Fail as e:
            out.append({"id": t["id"], "name": t["name"], "level": t["level"], "pass": False, "detail": str(e)})
        except (OSError, http.client.HTTPException) as e:
            out.append({"id": t["id"], "name": t["name"], "level": t["level"], "pass": False, "detail": "transport: %s" % e})
    return out


def requirements_md(report):
    lines = ["# Cache configuration: %s" % report["cache_label"], "",
             "Package version: %s. Configuration SHA-256: %s." % (report["package_version"], report["config_sha256"] or "none"), "",
             "| test | level | result | detail |", "|---|---|---|---|"]
    for r in report["results"]:
        res = ("pass" if r["pass"] else "FAIL") if r["level"] == "MUST" else (r["detail"] or "")
        lines.append("| %s %s | %s | %s | %s |" % (r["id"], r["name"], r["level"], res, "" if r["pass"] else r["detail"]))
    musts = [r for r in report["results"] if r["level"] == "MUST"]
    ok = all(r["pass"] for r in musts)
    lines += ["", "This configuration %s http-delivery §9 for every MUST test." % ("meets" if ok else "does not meet")]
    return "\n".join(lines) + "\n"


def main(argv):
    if len(argv) == 3 and argv[0] == "stage":
        stage(argv[1], argv[2])
        return 0
    if not argv or argv[0] != "run":
        print(__doc__)
        return 2
    opt = lambda f, d=None: argv[argv.index(f) + 1] if f in argv else d
    base, root = opt("--base"), opt("--root")
    if base is None or root is None:
        print(__doc__)
        return 2
    results = run_tests(base, root, opt("--origin"))
    cfg = opt("--config")
    report = {
        "cache_label": opt("--label", "unlabelled"),
        "package_version": opt("--version", "n/a"),
        "config_sha256": hashlib.sha256(open(cfg, "rb").read()).hexdigest() if cfg else None,
        "results": results,
    }
    if opt("--report"):
        open(opt("--report"), "w").write(json.dumps(report, indent=1) + "\n")
        open(os.path.join(os.path.dirname(opt("--report")) or ".", "requirements-%s.md" % report["cache_label"]), "w").write(requirements_md(report))
    failed = [r for r in results if r["level"] == "MUST" and not r["pass"]]
    for r in results:
        if not r["pass"]:
            print("FAIL %s: %s" % (r["id"], r["detail"]))
        elif r["level"] != "MUST":
            print("info %s: %s" % (r["id"], r["detail"]))
    print("%s: %d/%d MUST tests pass" % (report["cache_label"], sum(1 for r in results if r["level"] == "MUST" and r["pass"]),
                                         sum(1 for r in results if r["level"] == "MUST")))
    return 1 if failed else 0


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