"""Verify Air Priori data independently: signature, fingerprints, chain, timestamp and raw re-derivation.

Stand-alone: it shares no code with the Air Priori service, so a council officer or an auditor can read
every line of what it checks. For each UTC day in the range it reports PASS or FAIL on:

  1. signature    the download is signed by Air Priori's published Ed25519 key (X-AP-Signature)
  2. rows         each site's rows for the day hash to the fingerprint sealed in the day's manifest,
                  and no sealed site-day is missing from the download
  3. merkle       the manifest's Merkle root is the root of its leaves (RFC 6962 hashing)
  4. chain        each manifest names the SHA-256 of the one before it, up to a timestamped day
  5. timestamp    an independent RFC 3161 authority timestamped that day's manifest (the token's
                  hash matches; with OpenSSL installed, the authority's signature is checked too)
  6. raw          for a sample of site-days, the raw manufacturer responses are fetched and the published
                  transform re-run: the result must equal the sealed rows exactly

Requires Python 3 with:  pip install requests pandas cryptography rfc3161ng
OpenSSL (optional, for step 5's signature check) ships with Git for Windows, macOS and Linux.

    python verify_air_priori_data.py --key YOUR-KEY --from 2025-11-01 --to 2025-12-01
    python verify_air_priori_data.py --key ap_demo_lewisham2026 --from 2025-12-01 --to 2025-12-08 --sites LBL-01,LBL-02
"""
import argparse
import base64
import hashlib
import io
import json
import os
import random
import shutil
import subprocess
import sys
import tempfile

import pandas as pd
import requests

# The published transform (normalise.py): AirGradient field -> pollutant, native unit, ppb -> ug/m3 factor.
FIELDS = {"pm02_corrected": ("pm25", "ug/m3", None), "pm10_corrected": ("pm10", "ug/m3", None),
          "pm01_corrected": ("pm1", "ug/m3", None), "no2": ("no2", "ppb", 1.91), "o3": ("o3", "ppb", 2.00)}


def sha(b):
    return hashlib.sha256(b).hexdigest()


def cjson(o):
    return (json.dumps(o, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + "\n").encode("utf-8")


def canonical(df):
    """site,timestamp,pollutant,value_ugm3,value_native,unit - values to 4 dp - sorted by pollutant, time."""
    d = df.assign(_t=pd.to_datetime(df["timestamp_utc"], utc=True)).sort_values(["pollutant", "_t"])
    ts = d["_t"].dt.strftime("%Y-%m-%dT%H:%M:%SZ")
    lines = (d["site_id"].astype(str) + "," + ts + "," + d["pollutant"].astype(str) + ","
             + d["value_ugm3"].astype(float).map("{:.4f}".format) + ","
             + d["value_native"].astype(float).map("{:.4f}".format) + "," + d["unit_native"].astype(str))
    return "".join(x + "\n" for x in lines).encode("utf-8")


def leaf_hash(entry):
    return hashlib.sha256(b"\x00" + cjson(entry)).digest()


def mth(hs):
    if not hs:
        return hashlib.sha256(b"").digest()
    if len(hs) == 1:
        return hs[0]
    k = 1
    while k * 2 < len(hs):
        k *= 2
    return hashlib.sha256(b"\x01" + mth(hs[:k]) + mth(hs[k:])).digest()


def transform(body, site_id):
    """Raw manufacturer JSON -> rows, exactly as the published transform does it."""
    rows = []
    for m in json.loads(body):
        for field, (pol, unit, factor) in FIELDS.items():
            v = m.get(field)
            if v is None:
                continue
            v = float(v)
            rows.append({"site_id": site_id, "timestamp_utc": m["timestamp"], "pollutant": pol,
                         "value_ugm3": round(v * factor, 2) if factor else round(v, 2), "value_native": v,
                         "unit_native": unit})
    df = pd.DataFrame(rows)
    if len(df):
        df["timestamp_utc"] = pd.to_datetime(df["timestamp_utc"], utc=True)
        df = df.drop_duplicates(["site_id", "timestamp_utc", "pollutant"], keep="last")
    return df


def check_signature(pub_b64, body, headers):
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
    digest = sha(body)
    if headers.get("X-AP-Body-SHA256") != digest:
        return False, "body hash differs from the X-AP-Body-SHA256 header"
    try:
        Ed25519PublicKey.from_public_bytes(base64.b64decode(pub_b64)).verify(
            base64.b64decode(headers["X-AP-Signature"]), b"airpriori-v1:" + digest.encode())
    except Exception as e:                                  # noqa: BLE001
        return False, "signature does not verify (%s)" % type(e).__name__
    return True, "signed by key %s" % headers.get("X-AP-Key-Id")


def check_tsr(tsr, manifest_bytes, certs):
    import rfc3161ng
    resp = rfc3161ng.decode_timestamp_response(tsr)
    if int(resp["status"]["status"]) not in (0, 1):
        return False, "timestamp authority did not grant the token"
    tst = resp["timeStampToken"]
    imprint = bytes(tst.tst_info["messageImprint"]["hashedMessage"]).hex()
    when = rfc3161ng.get_timestamp(tst, naive=False).isoformat()
    if imprint != sha(manifest_bytes):
        return False, "the token covers a different file"
    exe = shutil.which("openssl")
    if not exe or not certs:
        return True, "token matches, issued %s (install OpenSSL to check the authority's signature too)" % when
    with tempfile.TemporaryDirectory() as d:
        paths = {}
        for name, data in [("t.tsr", tsr), ("m.json", manifest_bytes)] + list(certs.items()):
            paths[name] = os.path.join(d, name)
            with open(paths[name], "wb") as f:
                f.write(data)
        r = subprocess.run([exe, "ts", "-verify", "-in", paths["t.tsr"], "-data", paths["m.json"],
                            "-CAfile", paths["cacert.pem"], "-untrusted", paths["tsa.crt"]],
                           capture_output=True, text=True)
    ok = "Verification: OK" in (r.stdout + r.stderr)
    return ok, ("token matches and OpenSSL verifies the authority's signature; issued %s" % when) if ok \
        else "OpenSSL could not verify the authority's signature"


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--base", default="https://api.airpriori.com")
    ap.add_argument("--key", default=os.environ.get("AP_API_KEY"), required=os.environ.get("AP_API_KEY") is None)
    ap.add_argument("--from", dest="start", required=True, help="first UTC day, YYYY-MM-DD")
    ap.add_argument("--to", dest="end", required=True, help="day after the last, YYYY-MM-DD")
    ap.add_argument("--sites", default="all")
    ap.add_argument("--raw-sample", type=int, default=3, help="site-days to re-derive from raw (0 = none)")
    ap.add_argument("--expect-key-id", help="fail unless the signing key has this id (from Air Priori)")
    a = ap.parse_args(argv)
    base, h = a.base.rstrip("/"), {"X-API-Key": a.key}
    s = requests.Session()
    s.headers.update(h)
    fails = []

    def report(name, ok, msg):
        print("  %-10s %s  %s" % (name, "PASS" if ok else "FAIL", msg))
        if not ok:
            fails.append(name)

    key = s.get(base + "/v1/provenance/public-key", timeout=60).json()
    print("Air Priori signing key: %s (%s)" % (key["key_id"], key["algorithm"]))
    if a.expect_key_id and key["key_id"] != a.expect_key_id:
        print("FAIL: expected key %s" % a.expect_key_id)
        return 1

    # 1. one signed download: every pollutant, raw resolution, whole UTC days
    r = s.get(base + "/v1/measurements", params={"sites": a.sites, "pollutants": "all", "from": a.start,
                                                 "to": a.end, "resolution": "raw", "format": "csv"}, timeout=900)
    r.raise_for_status()
    print("Downloaded %s rows (request %s, dataset %s)" % (r.headers.get("X-Total-Rows"),
                                                         r.headers.get("X-Request-ID"), r.headers.get("X-Dataset")))
    print("\n[1] Signature")
    report("signature", *check_signature(key["public_key_b64"], r.content, r.headers))
    df = pd.read_csv(io.BytesIO(r.content))
    df["timestamp_utc"] = pd.to_datetime(df["timestamp_utc"], utc=True)
    df["day"] = df["timestamp_utc"].dt.strftime("%Y-%m-%d")
    have = {k: sha(canonical(g)) for k, g in df.groupby(["day", "site_id"])}
    counts = df.groupby(["day", "site_id"]).size().to_dict()

    # 2-5. the manifests, the chain and the timestamp
    m = s.get(base + "/v1/provenance/manifests", params={"from": a.start, "to": a.end}, timeout=300)
    m.raise_for_status()
    m = m.json()
    certs = {}
    for name in ("cacert.pem", "tsa.crt"):
        c = s.get(base + "/v1/provenance/tsa/" + name, timeout=60)
        if c.status_code == 200:
            certs[name] = c.content
    wanted = None if a.sites == "all" else set(x.strip() for x in a.sites.split(","))
    print("\n[2-3] Fingerprints and Merkle roots, day by day")
    prev = m["prev_manifest_sha256_before_first"]
    chain_ok, leaves_index = True, {}
    for day in m["days"]:
        b = day["manifest"].encode("utf-8")
        man = json.loads(b)
        if man["prev_manifest_sha256"] != prev:
            chain_ok = False
            report("chain", False, "%s does not follow the previous day" % day["date"])
        prev = sha(b)
        if day["date"] >= a.end:
            continue                                        # beyond the range: only a link to the anchor
        bad = []
        if mth([leaf_hash(e) for e in man["leaves"]]).hex() != man["merkle_root"]:
            bad.append("Merkle root")
        for e in man["leaves"]:
            if wanted and e["site_id"] not in wanted:
                continue
            leaves_index[(day["date"], e["site_id"])] = e
            if have.get((day["date"], e["site_id"])) != e["rows_sha256"]:
                bad.append("%s (%s rows sealed, %s downloaded)" % (e["site_id"], e["n_rows"],
                                                                   counts.get((day["date"], e["site_id"]), 0)))
        extra = sorted(st for (d, st) in have if d == day["date"] and (d, st) not in leaves_index)
        bad += ["%s not sealed" % x for x in extra]
        report(day["date"], not bad, "%d site(s) match the sealed fingerprints" % len(man["leaves"])
               if not bad else "mismatch: " + "; ".join(bad))
    print("\n[4] Chain")
    report("chain", chain_ok, "%d manifest(s) linked, %s to %s" % (len(m["days"]), m["days"][0]["date"],
                                                                     m["days"][-1]["date"]))
    print("\n[5] Independent timestamp")
    anchor = next((d for d in m["days"] if d["date"] == m["anchor"]), None)
    if anchor and anchor["tsr_b64"]:
        ok, msg = check_tsr(base64.b64decode(anchor["tsr_b64"]), anchor["manifest"].encode("utf-8"), certs)
        report("timestamp", ok, "%s (covers every day before it through the chain): %s" % (anchor["date"], msg))
    else:
        report("timestamp", False, "no timestamped day yet at or after the range")

    # 6. re-derive a sample of site-days from the raw manufacturer responses
    print("\n[6] Re-derivation from raw manufacturer data")
    sample = random.sample(sorted(leaves_index), min(a.raw_sample, len(leaves_index)))
    for day, site in sample:
        e = leaves_index[(day, site)]
        parts = []
        for digest in e["raw_sha256"]:
            rr = s.get(base + "/v1/raw/" + digest, timeout=120)
            if rr.status_code != 200 or sha(rr.content) != digest:
                parts = None
                break
            parts.append(transform(rr.content, site))
        if not parts:
            report("raw", False, "%s %s: raw response missing or altered" % (day, site))
            continue
        d = pd.concat(parts).drop_duplicates(["site_id", "timestamp_utc", "pollutant"], keep="last")
        d = d[d["timestamp_utc"].dt.strftime("%Y-%m-%d") == day]
        ok = sha(canonical(d)) == e["rows_sha256"]
        report("raw", ok, "%s %s: %d raw response(s) -> %d rows %s the sealed fingerprint"
               % (day, site, len(e["raw_sha256"]), len(d), "reproduce" if ok else "DO NOT reproduce"))

    print("\nRESULT: %s" % ("PASS - every check passed" if not fails else "FAIL - %d check(s) failed" % len(fails)))
    return 0 if not fails else 1


if __name__ == "__main__":
    sys.exit(main())