#!/usr/bin/env python3
"""
TIG Decentralised Verification — Stage 1: find candidate nonces to verify.

Given a round, challenge, and Benchmarker, this samples up to N nonces to verify
from that Benchmarker's benchmarks, never taking more than `num_bundles` nonces
from any single benchmark (beyond that, extra reports on one benchmark have no
consequence). Already-reported nonces are skipped, since they cannot be reported
again.

Algorithm, track, and compute type are optional filters — omit them to sample
across everything. Each candidate is labelled with all three, so you can group the
output afterwards.

It keeps the API footprint small:
  * GET /get-reports?round=<round>                                  (once, for dedup)
  * GET /get-reportable-benchmark-ids?round&player_id&challenge_id  (once)
  * GET /get-benchmark-data?benchmark_id=<id>   (once per benchmark of that Benchmarker)

Usage:
    pip install requests
    python3 find_candidate_nonces.py --round <round> --challenge c001 \
        --benchmarker 0x... -n 20 \
        [--algorithm c001_a098] [--track <track_id>] [--compute-type aws_c7i aws_m7i] [--testnet]

Output: one `<benchmark_id> <nonce> <compute_type> <track>` per line — each a
candidate to verify in Stage 2, labelled with the compute type and track.
"""
import argparse
import random

import requests

MAINNET_API = "https://mainnet-api.tig.foundation"
TESTNET_API = "https://testnet-api.tig.foundation"


def main():
    parser = argparse.ArgumentParser(description="Find candidate nonces to verify.")
    parser.add_argument("--round", type=int, required=True,
                        help="Round to audit (current or previous — the reporting window)")
    parser.add_argument("--challenge", required=True, help="Challenge id, e.g. c001")
    parser.add_argument("--benchmarker", required=True, help="Benchmarker (player) id, e.g. 0x...")
    parser.add_argument("-n", type=int, required=True,
                        help="Max number of nonces to sample. If fewer are available after "
                             "filtering, all of them are returned.")
    parser.add_argument("--algorithm", default=None,
                        help="Only sample benchmarks using this algorithm (e.g. c001_a098). "
                             "Default: any algorithm.")
    parser.add_argument("--track", default=None,
                        help="Only sample benchmarks on this track. Default: any track.")
    parser.add_argument("--compute-type", nargs="+", default=None,
                        help="Only sample nonces whose benchmark committed to one of these compute "
                             "types (e.g. aws_c7i aws_m7i). Default: any compute type.")
    parser.add_argument("--testnet", action="store_true", help="Use the testnet API")
    args = parser.parse_args()

    api = TESTNET_API if args.testnet else MAINNET_API
    compute_types = set(args.compute_type) if args.compute_type else None

    # Nonces already reported this round — one report per (benchmark, nonce), so skip these.
    reports = requests.get(f"{api}/get-reports", params={"round": args.round}).json()["reports"]
    already_reported = {(r["details"]["benchmark_id"], r["details"]["nonce"]) for r in reports}

    # This Benchmarker's benchmarks for this challenge in this round.
    benchmark_ids = requests.get(
        f"{api}/get-reportable-benchmark-ids",
        params={"round": args.round, "player_id": args.benchmarker, "challenge_id": args.challenge},
    ).json()["benchmark_ids"]

    # Keep those matching the optional algorithm / track / compute-type filters.
    benchmarks = []
    for benchmark_id in benchmark_ids:
        d = requests.get(f"{api}/get-benchmark-data", params={"benchmark_id": benchmark_id}).json()
        settings = d["precommit"]["settings"]
        details = d["precommit"]["details"]
        if args.algorithm and settings["algorithm_id"] != args.algorithm:
            continue
        if args.track and settings["track_id"] != args.track:
            continue
        if compute_types and details["compute_type"] not in compute_types:
            continue
        benchmarks.append({
            "benchmark_id": benchmark_id,
            "num_nonces": details["num_nonces"],
            "num_bundles": details["num_bundles"],  # cap: never sample more than this per benchmark
            "compute_type": details["compute_type"],
            "track": settings["track_id"],
        })

    if not benchmarks:
        print("No matching benchmarks found.")
        return

    # Build the pool of reportable nonces, then randomly sample up to N, taking at most
    # num_bundles from any one benchmark.
    pool = [
        (b, nonce)
        for b in benchmarks
        for nonce in range(b["num_nonces"])
        if (b["benchmark_id"], nonce) not in already_reported
    ]
    random.shuffle(pool)

    sample, per_benchmark = [], {}
    for b, nonce in pool:
        if len(sample) >= args.n:
            break
        if per_benchmark.get(b["benchmark_id"], 0) >= b["num_bundles"]:
            continue
        sample.append((b["benchmark_id"], nonce, b["compute_type"], b["track"]))
        per_benchmark[b["benchmark_id"]] = per_benchmark.get(b["benchmark_id"], 0) + 1

    for benchmark_id, nonce, compute_type, track in sample:
        print(benchmark_id, nonce, compute_type, track)


if __name__ == "__main__":
    main()