#!/usr/bin/env python3
"""
Independent recount and audit-chain verification.

This is a SECOND IMPLEMENTATION (brief section 4.4: "a second implementation or
spreadsheet/script should reproduce the totals from the anonymous export"). It
deliberately shares no code, no library and no language with the application
that produced the result. It reads only the sealed package, which contains no
voter identifiers at all.

An officer who does not trust the application can run this on the exported
files and compare. If the totals disagree, do not publish.

Usage
-----
    python3 tools/recount.py var/export/<election-id>
    python3 tools/recount.py var/export/<election-id> --audit audit.json

Exit status is 0 when everything reconciles and 1 when it does not, so it can
be wired into a check as well as read by a person.

Requires only the Python 3 standard library.
"""

import argparse
import hashlib
import json
import os
import sys


def load(path):
    with open(path, "r", encoding="utf-8") as handle:
        return json.load(handle)


def recount(ballots_doc):
    """Tally approval-block votes. Written from the published counting rule,
    not by translating the application's code."""
    candidates = {c["id"]: c["fullName"] for c in ballots_doc["candidates"]}
    max_selections = ballots_doc["maxSelections"]

    tally = {cid: 0 for cid in candidates}
    seen = set()
    blank = 0
    problems = []

    for ballot in ballots_doc["ballots"]:
        bid = ballot["ballotId"]
        if bid in seen:
            problems.append("duplicate ballot id %s" % bid)
            continue
        seen.add(bid)

        picks = list(dict.fromkeys(ballot["selections"]))
        if len(picks) != len(ballot["selections"]):
            problems.append("ballot %s repeats a candidate" % bid)
        if len(picks) > max_selections:
            problems.append(
                "ballot %s selects %d, above the maximum of %d"
                % (bid, len(picks), max_selections)
            )
        if not picks:
            blank += 1
        for cid in picks:
            if cid not in tally:
                problems.append("ballot %s names unknown candidate %s" % (bid, cid))
                continue
            tally[cid] += 1

    ordered = sorted(tally.items(), key=lambda kv: (-kv[1], kv[0]))
    return {
        "ballots_counted": len(seen),
        "blank": blank,
        "ordered": [(candidates[cid], cid, votes) for cid, votes in ordered],
        "tally": tally,
        "problems": problems,
    }


# The fractions a constitution names, as exact integers. Comparing
# votes/cast against 0.6667 in floating point decides a close two-thirds
# vote wrongly; n * denominator >= cast * numerator never rounds.
PASS_FRACTIONS = {
    "majority": (1, 2, True),          # strictly more than half
    "two_thirds": (2, 3, False),       # at least two thirds
    "three_quarters": (3, 4, False),
}

PASS_LABELS = {
    "majority": "more than half of the votes cast",
    "two_thirds": "at least two thirds of the votes cast",
    "three_quarters": "at least three quarters of the votes cast",
}


def decide_motion(ballots_doc, mine):
    """Was the motion carried? Written from the published rule, as above.

    The option marked inFavour is measured against the votes cast; every other
    option counts towards that total; abstentions are excluded from it.
    """
    favoured = [c for c in ballots_doc["candidates"] if c.get("inFavour")]
    if len(favoured) != 1:
        return None, ["a motion must have exactly one option marked as carrying it"]

    rule = ballots_doc.get("passRule", "majority")
    if rule not in PASS_FRACTIONS:
        return None, ["unknown pass rule %r" % rule]

    in_favour = mine["tally"].get(favoured[0]["id"], 0)
    cast = sum(mine["tally"].values())
    numerator, denominator, strict = PASS_FRACTIONS[rule]
    left, right = in_favour * denominator, cast * numerator
    carried = bool(cast) and (left > right if strict else left >= right)
    return {
        "carried": carried,
        "rule": rule,
        "in_favour": in_favour,
        "against": cast - in_favour,
        "cast": cast,
        "abstentions": mine["blank"],
    }, []


def check_manifest(directory):
    """Recompute the SHA-256 of every file the manifest names."""
    manifest_path = os.path.join(directory, "MANIFEST.sha256")
    if not os.path.exists(manifest_path):
        return None, ["no MANIFEST.sha256 in %s" % directory]

    problems = []
    checked = 0
    package_hash = None
    with open(manifest_path, "r", encoding="utf-8") as handle:
        for line in handle:
            line = line.rstrip("\n")
            if line.startswith("# package:"):
                package_hash = line.split(":", 1)[1].strip()
                continue
            if not line.strip():
                continue
            expected, name = line.split("  ", 1)
            target = os.path.join(directory, name)
            if not os.path.exists(target):
                problems.append("%s is named in the manifest but missing" % name)
                continue
            with open(target, "rb") as f:
                actual = hashlib.sha256(f.read()).hexdigest()
            if actual != expected:
                problems.append(
                    "%s does not match its manifest hash (file has been altered)" % name
                )
            checked += 1
    return {"checked": checked, "package_hash": package_hash}, problems


def canonical_payload(value):
    """Sort object keys at every depth, matching the application's canonical
    form so the same bytes are hashed."""
    if isinstance(value, list):
        return [canonical_payload(v) for v in value]
    if isinstance(value, dict):
        return {k: canonical_payload(value[k]) for k in sorted(value)}
    return value


def verify_audit_chain(audit_doc):
    """Replay the hash chain from genesis and recompute every event hash."""
    genesis = "0" * 64
    expected_prev = genesis
    head = genesis
    problems = []

    events = audit_doc["events"]
    for event in events:
        if event["prev_hash"] != expected_prev:
            problems.append(
                "event %s does not follow its predecessor (an event was removed or reordered)"
                % event["seq"]
            )
            break
        canonical = json.dumps(
            [
                str(event["seq"]),
                event["election_id"],
                event["occurred_at"],
                event["actor"],
                event["action"],
                event["object_type"],
                event["object_id"],
                event["reason"],
                event["approval_ref"],
                canonical_payload(event["payload"]),
                event["prev_hash"],
            ],
            separators=(",", ":"),
            ensure_ascii=False,
        )
        recomputed = hashlib.sha256(canonical.encode("utf-8")).hexdigest()
        if recomputed != event["hash"]:
            problems.append(
                "event %s does not match its recorded hash (it has been altered)"
                % event["seq"]
            )
            break
        expected_prev = event["hash"]
        head = event["hash"]

    return {"events": len(events), "head": head}, problems


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("directory", help="sealed package directory")
    parser.add_argument(
        "--audit",
        help="audit export JSON (from /admin/elections/<id>/audit.json) to verify as well",
    )
    args = parser.parse_args()

    failures = []
    # Kept apart from failures on purpose. A tie at the seat boundary is not a
    # defect in the count -- the arithmetic is complete and correct -- it is a
    # question the count is not allowed to answer, which the officers answer
    # under a rule agreed before any names were known. Publication is already
    # blocked in the application until they record what they did, so counting
    # it as a failure here told a member checking a properly published result
    # "do not publish" about something that had been published weeks earlier.
    decisions = []

    print("Sealed package: %s" % args.directory)
    print("=" * 66)

    manifest, manifest_problems = check_manifest(args.directory)
    if manifest:
        print("Manifest      : %d files, all hashes match" % manifest["checked"]
              if not manifest_problems else "Manifest      : PROBLEMS")
        if manifest["package_hash"]:
            print("Package hash  : %s" % manifest["package_hash"])
    for p in manifest_problems:
        print("  ! %s" % p)
        failures.append(p)

    ballots_doc = load(os.path.join(args.directory, "ballots.json"))
    mine = recount(ballots_doc)

    print("")
    print("Independent recount")
    print("-" * 66)
    print("Ballots counted: %d  (blank: %d)" % (mine["ballots_counted"], mine["blank"]))
    print("")
    width = max([len(n) for n, _, _ in mine["ordered"]] + [9])
    seats = ballots_doc["seats"]

    # A motion is decided by a threshold, not by filling seats, so it is
    # reported as a verdict and the seat markers below are skipped entirely.
    is_motion = ballots_doc.get("kind") == "motion"

    if is_motion:
        for name, _cid, votes in mine["ordered"]:
            print("  %-*s  %5d" % (width, name, votes))
    else:
        # A tie at the seat boundary leaves the seats inside it undecided; the
        # ones above it are still decided by the count.
        boundary_votes = None
        if (
            len(mine["ordered"]) > seats
            and mine["ordered"][seats - 1][2] == mine["ordered"][seats][2]
        ):
            boundary_votes = mine["ordered"][seats - 1][2]
        for index, (name, _cid, votes) in enumerate(mine["ordered"]):
            if boundary_votes is not None and votes == boundary_votes:
                marker = "  <- tied for the last seat"
            elif boundary_votes is not None:
                marker = "  <- fills a seat" if votes > boundary_votes else ""
            else:
                marker = "  <- fills a seat" if index < seats else ""
            print("  %-*s  %5d%s" % (width, name, votes, marker))

    for p in mine["problems"]:
        print("  ! %s" % p)
        failures.append(p)

    if is_motion:
        verdict, verdict_problems = decide_motion(ballots_doc, mine)
        for p in verdict_problems:
            print("  ! %s" % p)
            failures.append(p)
        if verdict:
            print("")
            print("  Decision : %s" % ("CARRIED" if verdict["carried"] else "NOT CARRIED"))
            print("  Needed   : %s" % PASS_LABELS[verdict["rule"]])
            print(
                "  In favour: %d of %d votes cast  (%d abstained, not counted)"
                % (verdict["in_favour"], verdict["cast"], verdict["abstentions"])
            )
            theirs_path = os.path.join(args.directory, "result.json")
            if os.path.exists(theirs_path):
                theirs = load(theirs_path).get("motion")
                if theirs is None:
                    msg = "result.json records no decision for a motion"
                    print("  ! %s" % msg)
                    failures.append(msg)
                elif bool(theirs.get("carried")) != verdict["carried"]:
                    msg = "DISAGREE: the application says %s, this recount says %s" % (
                        "carried" if theirs.get("carried") else "not carried",
                        "carried" if verdict["carried"] else "not carried",
                    )
                    print("  ! %s" % msg)
                    failures.append(msg)
                else:
                    print("  AGREE: the application reached the same decision.")
    elif len(mine["ordered"]) > seats:
        # Tie at the seat boundary: report it, never resolve it here.
        if mine["ordered"][seats - 1][2] == mine["ordered"][seats][2]:
            msg = (
                "the seat boundary fell inside a tie at %d votes. The rule agreed "
                "beforehand was: %s\n    What the officers did about it is recorded "
                "on the published result page and in the audit log." % (
                    mine["ordered"][seats - 1][2],
                    ballots_doc.get("tieRule", "(not recorded)"),
                )
            )
            print("  * %s" % msg)
            decisions.append(msg)

    result_path = os.path.join(args.directory, "result.json")
    if os.path.exists(result_path):
        theirs = load(result_path)
        their_totals = {t["candidateId"]: t["votes"] for t in theirs["totals"]}
        my_totals = {cid: votes for _n, cid, votes in mine["ordered"]}
        print("")
        print("Comparison with the application's result.json")
        print("-" * 66)
        if their_totals == my_totals:
            print("  AGREE: every candidate total matches.")
        else:
            for cid in sorted(set(their_totals) | set(my_totals)):
                if their_totals.get(cid) != my_totals.get(cid):
                    msg = "DISAGREE on %s: application %s, recount %s" % (
                        cid,
                        their_totals.get(cid),
                        my_totals.get(cid),
                    )
                    print("  ! %s" % msg)
                    failures.append(msg)

    summary_path = os.path.join(args.directory, "summary.json")
    if os.path.exists(summary_path):
        summary = load(summary_path)
        print("")
        print("Reconciliation")
        print("-" * 66)
        accepted = summary["acceptedBallots"]
        redeemed = summary["redeemedCredentials"]
        print("Accepted ballots     : %d" % accepted)
        print("Redeemed credentials : %d" % redeemed)
        if accepted != redeemed:
            msg = "accepted ballots (%d) != redeemed credentials (%d)" % (accepted, redeemed)
            print("  ! %s" % msg)
            failures.append(msg)
        if mine["ballots_counted"] != accepted:
            msg = "recount saw %d ballots, the summary claims %d" % (
                mine["ballots_counted"],
                accepted,
            )
            print("  ! %s" % msg)
            failures.append(msg)
        if not summary.get("configUnchangedSinceFreeze", True):
            msg = "the configuration changed after it was frozen"
            print("  ! %s" % msg)
            failures.append(msg)

    receipts_path = os.path.join(args.directory, "receipts.txt")
    if os.path.exists(receipts_path):
        with open(receipts_path, "r", encoding="utf-8") as handle:
            receipts = [line.strip() for line in handle if line.strip()]
        if len(receipts) != len(set(receipts)):
            msg = "the published receipt list contains duplicates"
            print("  ! %s" % msg)
            failures.append(msg)
        if len(receipts) != mine["ballots_counted"]:
            msg = "%d receipts published but %d ballots counted" % (
                len(receipts),
                mine["ballots_counted"],
            )
            print("  ! %s" % msg)
            failures.append(msg)

    if args.audit:
        audit_doc = load(args.audit)
        chain, chain_problems = verify_audit_chain(audit_doc)
        print("")
        print("Audit chain")
        print("-" * 66)
        print("Events replayed: %d" % chain["events"])
        print("Head hash      : %s" % chain["head"])
        if chain_problems:
            for p in chain_problems:
                print("  ! %s" % p)
                failures.append(p)
        else:
            print("  Every event hash recomputed and matched.")

    print("")
    print("=" * 66)
    if failures:
        print(
            "RESULT: %d problem(s), marked ! above. The published totals are not "
            "supported by the published ballots, or a file has been altered. Ask "
            "the officers to explain before relying on this result." % len(failures)
        )
        return 1
    if decisions:
        print("RESULT: the count reconciles.")
        print(
            "        %d matter(s), marked * above, needed the officers' decision "
            "rather than the count's. That is not a fault: check what they\n"
            "        recorded against the rule agreed in advance." % len(decisions)
        )
        return 0
    print("RESULT: everything reconciles.")
    return 0


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