#!/usr/bin/env python3
"""
Draw-select independent verifier (Python 3, stdlib only).

Replicates the full draw algorithm from a draw-select certificate:
  1. Commitment hash  — SHA-256(drawID_UTF8 | totalEntries_BE32 | numberOfWinners_BE32 | numberOfAlternates_BE32 | revealSeed)
  2. Deterministic RNG — counter-mode SHA-256 with decimal-ASCII counter
  3. Transcript        — binary append-only log
  4. Winners           — rejection-sampling without-replacement selection (no duplicates)
  5. Alternates        — drawn from remaining pool after winners, same RNG stream

Usage:
    python3 verify.py \\
        --draw-id "my-draw-2024" \\
        --total-entries 1000 \\
        --number-of-winners 3 \\
        --reveal-seed aabbcc... \\
        [--number-of-alternates 2] \\
        [--commitment-hash ddeeff...] \\
        [--transcript-hash 112233...] \\
        [--winners 42,7,999] \\
        [--alternates 5,18]

Exit codes:
    0 — verification passed (or computed without values to check against)
    1 — verification failed (hash or winner/alternate mismatch)
    2 — usage / input error
"""

import argparse
import hashlib
import struct
import sys


# ── SeededRNG ─────────────────────────────────────────────────────────────────

class SeededRNG:
    """Counter-mode SHA-256 RNG.  Counter encoded as decimal ASCII string."""
    def __init__(self, seed: bytes):
        self._seed = seed
        self._counter = 0

    def uint32(self) -> int:
        h = hashlib.sha256()
        h.update(self._seed)
        h.update(str(self._counter).encode('ascii'))
        self._counter += 1
        return struct.unpack('>I', h.digest()[:4])[0]


# ── Transcript ────────────────────────────────────────────────────────────────

class Transcript:
    """Binary append-only draw transcript.

    Record types:
        0x01 + uint32BE  — RNG sample or accepted value
        0x02 + uint32BE  — winner index
        0x03 + uint32BE  — alternate index
        0x10 + uint32BE(len) + UTF-8 bytes — string (draw_id at start)
        0xff             — rejected sample (bias-prevention discard)
    """
    def __init__(self):
        self._chunks: list[bytes] = []

    def push_uint32(self, v: int) -> None:
        self._chunks.append(b'\x01' + struct.pack('>I', v))

    def push_reject(self) -> None:
        self._chunks.append(b'\xff')

    def push_winner(self, idx: int) -> None:
        self._chunks.append(b'\x02' + struct.pack('>I', idx))

    def push_alternate(self, idx: int) -> None:
        self._chunks.append(b'\x03' + struct.pack('>I', idx))

    def push_string(self, s: str) -> None:
        b = s.encode('utf-8')
        self._chunks.append(b'\x10')
        self._chunks.append(struct.pack('>I', len(b)))
        self._chunks.append(b)

    def to_bytes(self) -> bytes:
        return b''.join(self._chunks)

    def hash_hex(self) -> str:
        return hashlib.sha256(self.to_bytes()).hexdigest()


def new_transcript(draw_id: str) -> Transcript:
    t = Transcript()
    t.push_string(draw_id)
    return t


# ── uniformIntWithTranscript ──────────────────────────────────────────────────

def uniform_int_with_transcript(max_val: int, rng: SeededRNG, t: Transcript):
    """Rejection-sampling uniform integer in [0, max_val).
    Records every sample in the transcript so the full execution can be audited.
    Returns (value, discarded_count).
    """
    if max_val == 1:
        t.push_uint32(0)
        t.push_uint32(0)
        return 0, 0

    limit = (2**32 // max_val) * max_val
    discarded = 0
    while True:
        r = rng.uint32()
        t.push_uint32(r)
        if r < limit:
            v = r % max_val
            t.push_uint32(v)
            return v, discarded
        t.push_reject()
        discarded += 1


# ── Commitment hash ───────────────────────────────────────────────────────────

def compute_commitment(draw_id: str, total_entries: int, number_of_winners: int,
                       number_of_alternates: int, seed: bytes) -> str:
    """SHA-256(drawID_UTF8 | uint32BE(totalEntries) | uint32BE(numberOfWinners) | uint32BE(numberOfAlternates) | seed)"""
    h = hashlib.sha256()
    h.update(draw_id.encode('utf-8'))
    h.update(struct.pack('>I', total_entries))
    h.update(struct.pack('>I', number_of_winners))
    h.update(struct.pack('>I', number_of_alternates))
    h.update(seed)
    return h.hexdigest()


# ── Main draw replication ─────────────────────────────────────────────────────

def run_draw(draw_id: str, total_entries: int, number_of_winners: int,
             number_of_alternates: int, seed: bytes):
    """Replicate the full draw.
    Returns (commitment_hex, transcript_hex, winners, alternates, discarded).
    """
    commitment_hex = compute_commitment(
        draw_id, total_entries, number_of_winners, number_of_alternates, seed
    )

    rng = SeededRNG(seed)
    transcript = new_transcript(draw_id)
    transcript.push_uint32(total_entries)
    transcript.push_uint32(number_of_winners)
    transcript.push_uint32(number_of_alternates)

    selected_set = set()
    total_discarded = 0

    winners = []
    while len(winners) < number_of_winners:
        v, d = uniform_int_with_transcript(total_entries, rng, transcript)
        total_discarded += d
        if v not in selected_set:
            selected_set.add(v)
            winners.append(v)
            transcript.push_winner(v)

    alternates = []
    while len(alternates) < number_of_alternates:
        v, d = uniform_int_with_transcript(total_entries, rng, transcript)
        total_discarded += d
        if v not in selected_set:
            selected_set.add(v)
            alternates.append(v)
            transcript.push_alternate(v)

    transcript_hex = transcript.hash_hex()
    return commitment_hex, transcript_hex, winners, alternates, total_discarded


# ── CLI ───────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description='Draw-select independent verifier',
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument('--draw-id', required=True,
                        help='draw_id from the certificate')
    parser.add_argument('--total-entries', type=int, required=True,
                        help='total_entries from the certificate')
    parser.add_argument('--number-of-winners', type=int, required=True,
                        help='number_of_winners from the certificate')
    parser.add_argument('--number-of-alternates', type=int, default=0,
                        help='number_of_alternates from the certificate (default: 0)')
    parser.add_argument('--reveal-seed', required=True,
                        help='reveal_seed (hex) from the certificate')
    parser.add_argument('--commitment-hash', default='',
                        help='commitment_hash to verify against (optional)')
    parser.add_argument('--transcript-hash', default='',
                        help='transcript_hash to verify against (optional)')
    parser.add_argument('--winners', default='',
                        help='comma-separated expected winners (optional)')
    parser.add_argument('--alternates', default='',
                        help='comma-separated expected alternates (optional)')
    parser.add_argument('--one-indexed', action='store_true',
                        help='winners/alternates on certificate are 1-indexed; pass them as-is and this flag handles the conversion')
    args = parser.parse_args()

    # Validate seed
    try:
        seed = bytes.fromhex(args.reveal_seed)
    except ValueError as e:
        print(f'verify: invalid --reveal-seed: {e}', file=sys.stderr)
        sys.exit(2)
    if len(seed) != 32:
        print(f'verify: --reveal-seed must be 32 bytes (64 hex chars), got {len(seed)}', file=sys.stderr)
        sys.exit(2)

    total = args.total_entries
    n = args.number_of_winners
    na = args.number_of_alternates
    if total <= 0:
        print('verify: --total-entries must be >= 1', file=sys.stderr)
        sys.exit(2)
    if n <= 0 or n > total:
        print('verify: --number-of-winners must be in [1, total_entries]', file=sys.stderr)
        sys.exit(2)
    if na < 0 or na > total - n:
        print('verify: --number-of-alternates must be in [0, total_entries - number_of_winners]', file=sys.stderr)
        sys.exit(2)

    commitment_hex, transcript_hex, winners, alternates, discarded = run_draw(
        args.draw_id, total, n, na, seed
    )

    display_winners = [w + 1 for w in winners] if args.one_indexed else winners
    display_alternates = [a + 1 for a in alternates] if args.one_indexed else alternates
    idx_label = '1-idx' if args.one_indexed else '0-idx'
    winners_str = ','.join(str(w) for w in display_winners)
    alternates_str = ','.join(str(a) for a in display_alternates)

    print('=== Draw Verification ===')
    print(f'Draw ID:           {args.draw_id}')
    print(f'Total Entries:     {total}')
    print(f'Number of Winners: {n}')
    if na > 0:
        print(f'Number of Alternates: {na}')
    print(f'Reveal Seed:       {args.reveal_seed}')
    print(f'Commitment Hash:   {commitment_hex}')
    print(f'Transcript Hash:   {transcript_hex}')
    print(f'Discarded Samples: {discarded}')
    print(f'Winners ({idx_label}):   {winners_str}')
    if na > 0:
        print(f'Alternates ({idx_label}): {alternates_str}')
    print()

    failed = False

    if args.commitment_hash:
        if args.commitment_hash == commitment_hex:
            print('✓ commitment_hash  MATCH')
        else:
            print('✗ commitment_hash  MISMATCH')
            print(f'  expected: {args.commitment_hash}')
            print(f'  computed: {commitment_hex}')
            failed = True

    if args.transcript_hash:
        if args.transcript_hash == transcript_hex:
            print('✓ transcript_hash  MATCH')
        else:
            print('✗ transcript_hash  MISMATCH')
            print(f'  expected: {args.transcript_hash}')
            print(f'  computed: {transcript_hex}')
            failed = True

    if args.winners:
        if args.winners == winners_str:
            print('✓ winners          MATCH')
        else:
            print('✗ winners          MISMATCH')
            print(f'  expected: {args.winners}')
            print(f'  computed: {winners_str}')
            failed = True

    if args.alternates:
        if args.alternates == alternates_str:
            print('✓ alternates       MATCH')
        else:
            print('✗ alternates       MISMATCH')
            print(f'  expected: {args.alternates}')
            print(f'  computed: {alternates_str}')
            failed = True

    if failed:
        print('\nVerification FAILED.')
        sys.exit(1)

    print('Verification PASSED.')
    sys.exit(0)


if __name__ == '__main__':
    main()
