#!/usr/bin/env python3
"""Verifiable random task draw for acceptance testing (acceptance-test.posttrain.dev).

The draw is deterministic given (task manifest, drand beacon round), and the round is
pinned to the time the vendor's manifest email was sent, so neither side can steer or
re-roll it. Anyone can re-run this script and get the same result.

Protocol (web-first):
  1. Vendor commits their dataset.toml at acceptance-test.posttrain.dev/dataset_draw.
     The commit time pins the beacon round (first drand round published after it),
     and the server returns a receipt naming the round, randomness, and drawn tasks.
  2. Anyone re-checks that receipt with:
       python3 verifiable_draw.py verify --manifest dataset.toml --receipt receipt.json
     All checks are recomputed independently, including fetching the beacon from drand.
  3. `draw` also exists as an offline fallback (email flow): commit by sending the
     manifest, then run:
       python3 verifiable_draw.py draw --manifest dataset.toml --sent 2026-08-26T14:03:00+09:00

Draw rule:
  M     = SHA256 of the sorted "sha256  task-name" manifest lines
  R     = first drand round published strictly AFTER the sent time
  B     = drand randomness of round R (public, verifiable, archived forever)
  seed  = SHA256(M || B)
  score = SHA256(seed || task_sha) per task; the n lowest scores are drawn.

Stdlib only. drand docs: https://drand.love
"""

import argparse
import hashlib
import json
import pathlib
import sys
import time
import urllib.request
from datetime import datetime, timezone

DRAND_BASE = "https://api.drand.sh"  # League of Entropy default chain (30s period)


def fetch_json(url: str) -> dict:
    with urllib.request.urlopen(url, timeout=30) as r:
        return json.loads(r.read())


def chain_info() -> dict:
    return fetch_json(f"{DRAND_BASE}/info")


def build_manifest(tasks_dir: str | None, manifest_file: str | None) -> list[tuple[str, str]]:
    """Return [(name, sha256_of_task_toml)], sorted by sha."""
    tasks: list[tuple[str, str]] = []
    if tasks_dir:
        root = pathlib.Path(tasks_dir)
        tomls = sorted(root.glob("*/task.toml"))
        if not tomls:
            sys.exit(f"no */task.toml found under {root}")
        for toml in tomls:
            tasks.append((toml.parent.name, hashlib.sha256(toml.read_bytes()).hexdigest()))
    elif manifest_file:
        import re
        text = pathlib.Path(manifest_file).read_text()
        toml_tasks = re.findall(r'name\s*=\s*"([^"]+)"\s*\n\s*digest\s*=\s*"sha256:([0-9a-fA-F]{64})"', text)
        if len(toml_tasks) >= 2:  # Harbor dataset.toml
            for name, sha in toml_tasks:
                tasks.append((name.split("/")[-1], sha.lower()))
        else:
            for line in text.splitlines():
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                sha, name = line.split(None, 1)
                name = name.strip().removeprefix("./")
                if name.endswith("/task.toml"):
                    name = name[: -len("/task.toml")]
                name = name.split("/")[-1] or name
                tasks.append((name, sha.lower()))
    else:
        sys.exit("provide --tasks-dir or --manifest")
    return sorted(tasks, key=lambda t: t[1])


def manifest_hash(tasks: list[tuple[str, str]]) -> str:
    text = "\n".join(f"{sha}  {name}" for name, sha in tasks)
    return hashlib.sha256(text.encode()).hexdigest()


def round_after(sent_iso: str, info: dict) -> tuple[int, int]:
    """First round published strictly after the sent time. Returns (round, publish_unix)."""
    t = datetime.fromisoformat(sent_iso.replace("Z", "+00:00"))
    if t.tzinfo is None:
        sys.exit("--sent must include a timezone, e.g. 2026-08-26T14:03:00+09:00")
    unix = t.timestamp()
    genesis, period = info["genesis_time"], info["period"]
    if unix < genesis:
        sys.exit("--sent is before the drand chain genesis")
    r = int((unix - genesis) // period) + 2  # round r publishes at genesis + (r-1)*period > sent
    return r, int(genesis + (r - 1) * period)


def compute_draw(tasks: list[tuple[str, str]], m_hash: str, randomness: str, n: int):
    seed = hashlib.sha256(bytes.fromhex(m_hash) + bytes.fromhex(randomness)).hexdigest()
    scored = sorted(
        tasks,
        key=lambda t: hashlib.sha256(bytes.fromhex(seed) + bytes.fromhex(t[1])).hexdigest(),
    )
    return seed, scored[:n]


def cmd_draw(args):
    tasks = build_manifest(args.tasks_dir, args.manifest)
    m_hash = manifest_hash(tasks)
    info = chain_info()
    rnd, publish_unix = round_after(args.sent, info)

    wait = publish_unix - time.time()
    if wait > 0:
        print(f"round {rnd} publishes in {int(wait) + 1}s, waiting...", file=sys.stderr)
        time.sleep(wait + 1)
    beacon = fetch_json(f"{DRAND_BASE}/public/{rnd}")

    seed, drawn = compute_draw(tasks, m_hash, beacon["randomness"], args.n)
    receipt = {
        "protocol": "acceptance-test.posttrain.dev verifiable draw v1",
        "task_count": len(tasks),
        "manifest_sha256": m_hash,
        "sent_time": args.sent,
        "drand_chain_hash": info["hash"],
        "drand_round": rnd,
        "drand_randomness": beacon["randomness"],
        "seed": seed,
        "n": args.n,
        "drawn": [{"name": name, "task_toml_sha256": sha} for name, sha in drawn],
    }
    out = pathlib.Path(args.out)
    out.write_text(json.dumps(receipt, indent=2) + "\n")
    print(f"manifest: {len(tasks)} tasks, M={m_hash[:16]}...")
    print(f"beacon:   round {rnd} (first after {args.sent})")
    print(f"drawn {args.n}:")
    for name, sha in drawn:
        print(f"  {sha[:12]}  {name}")
    print(f"\nreceipt written to {out} — email it back together with the drawn task folders.")


def cmd_verify(args):
    receipt = json.loads(pathlib.Path(args.receipt).read_text())
    tasks = build_manifest(args.tasks_dir, args.manifest)
    checks: list[tuple[str, bool, str]] = []

    m_hash = manifest_hash(tasks)
    checks.append(("manifest hash matches receipt", m_hash == receipt["manifest_sha256"],
                   f"local {m_hash[:16]} vs receipt {receipt['manifest_sha256'][:16]}"))

    info = chain_info()
    rnd, _ = round_after(receipt["sent_time"], info)
    checks.append(("round is first after sent_time", rnd == receipt["drand_round"],
                   f"expected {rnd}, receipt {receipt['drand_round']}"))

    beacon = fetch_json(f"{DRAND_BASE}/public/{receipt['drand_round']}")
    checks.append(("beacon randomness authentic (fetched independently)",
                   beacon["randomness"] == receipt["drand_randomness"],
                   f"drand {beacon['randomness'][:16]} vs receipt {receipt['drand_randomness'][:16]}"))

    seed, drawn = compute_draw(tasks, m_hash, beacon["randomness"], receipt["n"])
    expect = [(d["name"], d["task_toml_sha256"]) for d in receipt["drawn"]]
    checks.append(("recomputed draw matches receipt", drawn == expect,
                   "drawn set recomputed from manifest + beacon"))

    if args.delivered_dir:
        ok, detail = True, []
        for d in receipt["drawn"]:
            toml = pathlib.Path(args.delivered_dir) / d["name"] / "task.toml"
            if not toml.exists():
                ok, _ = False, detail.append(f"missing {d['name']}")
                continue
            if hashlib.sha256(toml.read_bytes()).hexdigest() != d["task_toml_sha256"]:
                ok, _ = False, detail.append(f"hash mismatch {d['name']}")
        checks.append(("delivered tasks match drawn SHAs", ok, "; ".join(detail) or "all match"))

    print("MANUAL: confirm sent_time is not earlier than the manifest email's arrival time.\n")
    failed = False
    for name, ok, detail in checks:
        print(f"  [{'PASS' if ok else 'FAIL'}] {name}  ({detail})")
        failed |= not ok
    sys.exit(1 if failed else 0)


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = p.add_subparsers(dest="cmd", required=True)

    d = sub.add_parser("draw", help="vendor: run the draw for your manifest")
    d.add_argument("--tasks-dir", help="directory containing <task>/task.toml folders")
    d.add_argument("--manifest", help="or a manifest file with '<sha256>  <name>' lines")
    d.add_argument("--sent", required=True, help="ISO time your manifest email was sent, e.g. 2026-08-26T14:03:00+09:00")
    d.add_argument("--n", type=int, default=10)
    d.add_argument("--out", default="receipt.json")
    d.set_defaults(func=cmd_draw)

    v = sub.add_parser("verify", help="buyer side: re-check a draw receipt")
    v.add_argument("--tasks-dir")
    v.add_argument("--manifest")
    v.add_argument("--receipt", required=True)
    v.add_argument("--delivered-dir", help="directory with the delivered drawn tasks to hash-check")
    v.set_defaults(func=cmd_verify)

    args = p.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()