#!/usr/bin/env python3
"""Reference implementation of the PCAT 702 verifiable draw.

Pure function: (pot list, rules, payments in canonical order) -> assignments,
refunds, entropy record.  Standard library only.  No network, no clock, no
randomness source other than the entropy carried inside each payment record.

Canonical order of payments (consensus-committed, see VERIFY.md):
    (block height, logical position inside the block)
Draws inside one payment are ordered by their index j = 0..n-1.

Seed of one payment:
    seed = SHA256( DOMAIN || lp(pot_fingerprint) || lp(txid) || lp(hash of block P+1) || lp(drand) )
lp(x) = 4-byte big-endian length of x followed by the raw bytes of x.
Both entropy inputs are mandatory under the published rules (Rules.require_*);
the relaxed mode exists only for experiments and is never used for the sale.

Draw j of a payment from a group with m items still free:
    u   = SHA256( lp(seed) || lp(group) || u32(j) || u32(ctr) )  as 256-bit int
    reject if u >= (2^256 // m) * m, retry with ctr+1       (no modulo bias)
    idx = u mod m ; item = free[idx] ; free.remove(idx)     (no replacement)
"""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from dataclasses import dataclass, field

DOMAIN = b"PCAT702/draw/v1"
UINT256 = 1 << 256

GROUP_POT = "pot"
GROUP_EPIC = "epic"
GROUP_ULTRA = "ultra"

# --------------------------------------------------------------------------
# Published parameters of the PCAT 702 sale (single source for draw.py, verify.py
# and the website).  Changing any of these changes every result.
# --------------------------------------------------------------------------

PCAT702_FINGERPRINT = "e6b640206edf7f58092d891210987152f4409f2f230d1a98b9f3eb76373ed364"
PCAT702_POT_COMPOSITION = {
    "Common": 414,
    "Uncommon": 191,
    "Rare": 69,
    "Epic": 15,
    "Ultra Rare": 1,
}
PCAT702_DIRECT_EPICS = 8
PCAT702_DIRECT_ULTRAS = 4

SEED_OFFSET = 1            # decisive block D = P + 1 (P = block of the payment)
FINALITY_DEPTH = 24        # draw only when D has 24 blocks on top of it (tip >= D + 24)
BEACON = "drand"           # public randomness mixed into every seed
BEACON_MARGIN_S = 1800     # drand round = first round at or after timestamp(D) + margin
                           # (must exceed the consensus reorg window of 17 blocks even with
                           #  back-dated timestamps and slow blocks: 17 x 30 s + 180 s < 1800 s)


class DrawError(ValueError):
    """Invalid input (never silently repaired)."""


# --------------------------------------------------------------------------
# Hash primitives
# --------------------------------------------------------------------------

def lp(data: bytes) -> bytes:
    """Length-prefix a field (u32 big endian) so concatenations are unambiguous."""
    return len(data).to_bytes(4, "big") + data


def canonical_json(obj) -> bytes:
    return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=True).encode("ascii")


def _hex32(value: str, what: str) -> bytes:
    try:
        raw = bytes.fromhex(value)
    except (ValueError, TypeError):
        raise DrawError(f"{what}: not hex: {value!r}") from None
    if len(raw) != 32:
        raise DrawError(f"{what}: expected 32 bytes, got {len(raw)}")
    return raw


def _stream(seed: bytes, group: str, j: int, ctr: int) -> int:
    digest = hashlib.sha256(lp(seed) + lp(group.encode("ascii")) + j.to_bytes(4, "big") + ctr.to_bytes(4, "big")).digest()
    return int.from_bytes(digest, "big")


def rand_below(seed: bytes, group: str, j: int, m: int) -> int:
    """Uniform integer in [0, m) from a 256-bit hash stream, rejection sampled."""
    if m < 1:
        raise DrawError("rand_below needs m >= 1")
    limit = (UINT256 // m) * m  # largest multiple of m that fits in 2^256
    ctr = 0
    while True:
        u = _stream(seed, group, j, ctr)
        if u < limit:
            return u % m
        ctr += 1


# --------------------------------------------------------------------------
# Pot list
# --------------------------------------------------------------------------

@dataclass(frozen=True)
class Item:
    id: str
    rarity: str
    meta_hash: str

    def as_dict(self) -> dict:
        return {"id": self.id, "rarity": self.rarity, "meta_hash": self.meta_hash}


@dataclass
class PotList:
    """Three ordered groups.  The order inside each group is part of the fingerprint."""

    pot: list
    epics: list
    ultras: list

    @staticmethod
    def from_json(doc: dict) -> "PotList":
        try:
            groups = {k: [Item(str(i["id"]), str(i["rarity"]), str(i["meta_hash"])) for i in doc[k]]
                      for k in ("pot", "direct_epic", "direct_ultra")}
        except (KeyError, TypeError) as exc:
            raise DrawError(f"pot list malformed: {exc!r}") from None
        lst = PotList(groups["pot"], groups["direct_epic"], groups["direct_ultra"])
        lst.validate()
        return lst

    def to_json(self) -> dict:
        return {
            "pot": [i.as_dict() for i in self.pot],
            "direct_epic": [i.as_dict() for i in self.epics],
            "direct_ultra": [i.as_dict() for i in self.ultras],
        }

    def validate(self) -> None:
        seen = set()
        for item in self.pot + self.epics + self.ultras:
            if item.id in seen:
                raise DrawError(f"duplicate NFT id {item.id!r}")
            seen.add(item.id)
            _hex32(item.meta_hash, f"meta_hash of {item.id}")

    def fingerprint(self) -> str:
        """SHA-256 over the canonical JSON of the complete list (hex).

        The published pot.json is exactly these canonical bytes, so
        `sha256sum pot.json` reproduces the fingerprint without this code."""
        return hashlib.sha256(canonical_json(self.to_json())).hexdigest()

    def check_pcat702(self) -> None:
        counts = {}
        for item in self.pot:
            counts[item.rarity] = counts.get(item.rarity, 0) + 1
        if counts != PCAT702_POT_COMPOSITION:
            raise DrawError(f"pot composition {counts} != {PCAT702_POT_COMPOSITION}")
        if len(self.epics) != PCAT702_DIRECT_EPICS or any(i.rarity != "Epic" for i in self.epics):
            raise DrawError("direct epics must be 8 x Epic")
        if len(self.ultras) != PCAT702_DIRECT_ULTRAS or any(i.rarity != "Ultra Rare" for i in self.ultras):
            raise DrawError("direct ultras must be 4 x Ultra Rare")


# --------------------------------------------------------------------------
# Rules and payments
# --------------------------------------------------------------------------

MICRO = 1_000_000


@dataclass(frozen=True)
class Rules:
    unit_price: int = 23 * MICRO     # per random cat from the pot
    max_cats: int = 10               # n = 1..max_cats
    epic_price: int = 223 * MICRO
    ultra_price: int = 523 * MICRO
    refund_fee: int = 9_000          # deducted from every refund (micro-NOID); covers the network fee
    partial_fill: bool = True        # pot has fewer than n cats: draw the rest, refund the missing ones
    require_seed_block: bool = True  # a drawn payment must carry the hash of block P+1
    require_beacon: bool = True      # ... and the public beacon value

    def __post_init__(self):
        for name in ("unit_price", "epic_price", "ultra_price"):
            v = getattr(self, name)
            if not isinstance(v, int) or isinstance(v, bool) or v <= 0:
                raise DrawError(f"rules: {name} must be a positive integer, got {v!r}")
        if not isinstance(self.max_cats, int) or isinstance(self.max_cats, bool) or self.max_cats < 1:
            raise DrawError(f"rules: max_cats must be >= 1, got {self.max_cats!r}")
        if not isinstance(self.refund_fee, int) or isinstance(self.refund_fee, bool) or self.refund_fee < 0:
            raise DrawError(f"rules: refund_fee must be >= 0, got {self.refund_fee!r}")
        if self.epic_price == self.ultra_price:
            raise DrawError("rules: epic_price and ultra_price collide")
        for n in range(1, self.max_cats + 1):
            if n * self.unit_price in (self.epic_price, self.ultra_price):
                raise DrawError(f"rules: {n} x unit_price collides with a direct price")
        for name in ("partial_fill", "require_seed_block", "require_beacon"):
            if not isinstance(getattr(self, name), bool):
                raise DrawError(f"rules: {name} must be true/false")

    def classify(self, amount: int):
        """Return (group, count) or (None, 0) when the amount is not a valid price."""
        if amount == self.epic_price:
            return GROUP_EPIC, 1
        if amount == self.ultra_price:
            return GROUP_ULTRA, 1
        if amount > 0 and amount % self.unit_price == 0 and 1 <= amount // self.unit_price <= self.max_cats:
            return GROUP_POT, amount // self.unit_price
        return None, 0

    def to_json(self) -> dict:
        return {k: getattr(self, k) for k in ("unit_price", "max_cats", "epic_price", "ultra_price", "refund_fee",
                                               "partial_fill", "require_seed_block", "require_beacon")}


@dataclass(frozen=True)
class Payment:
    height: int
    block_hash: str
    txid: str
    position: int
    amount: int                 # micro-NOID received by the founder address in this tx
    sender: str                 # input_owner
    seed_block_hash: str = ""   # hash of the decisive block P+1 ("" only in relaxed mode: own block)
    beacon: str = ""            # 32-byte public randomness (hex); "" only in relaxed mode
    beacon_round: int = 0       # informational: the beacon round the value came from (archived in the ledger)

    @staticmethod
    def from_json(d: dict) -> "Payment":
        try:
            p = Payment(
                height=int(d["height"]), block_hash=str(d["block_hash"]), txid=str(d["txid"]),
                position=int(d["position"]), amount=int(d["amount"]), sender=str(d["sender"]),
                seed_block_hash=str(d.get("seed_block_hash", "") or ""), beacon=str(d.get("beacon", "") or ""),
                beacon_round=int(d.get("beacon_round", 0) or 0),
            )
        except (KeyError, TypeError, ValueError) as exc:
            raise DrawError(f"payment malformed: {exc!r}") from None
        p.validate()
        return p

    def validate(self) -> None:
        _hex32(self.block_hash, "block_hash")
        _hex32(self.txid, "txid")
        if self.seed_block_hash:
            _hex32(self.seed_block_hash, "seed_block_hash")
        if self.beacon:
            _hex32(self.beacon, "beacon")
        if self.amount <= 0 or self.height < 0 or self.position < 0 or self.beacon_round < 0:
            raise DrawError(f"payment {self.txid}: non-positive amount or negative height/position/round")
        if not self.sender:
            raise DrawError(f"payment {self.txid}: sender missing")

    def to_json(self) -> dict:
        out = {"height": self.height, "block_hash": self.block_hash, "txid": self.txid,
               "position": self.position, "amount": self.amount, "sender": self.sender}
        if self.seed_block_hash:
            out["seed_block_hash"] = self.seed_block_hash
        if self.beacon:
            out["beacon"] = self.beacon
        if self.beacon_round:
            out["beacon_round"] = self.beacon_round
        return out

    def entropy(self) -> list:
        """Ordered entropy inputs of this payment."""
        parts = [_hex32(self.seed_block_hash or self.block_hash, "seed block")]
        if self.beacon:
            parts.append(_hex32(self.beacon, "beacon"))
        return parts

    def entropy_record(self) -> dict:
        """What the ledger archives so the draw stays reproducible without any external service."""
        return {"txid": self.txid, "height": self.height, "position": self.position, "block_hash": self.block_hash,
                "seed_block_hash": self.seed_block_hash or self.block_hash, "beacon": self.beacon,
                "beacon_round": self.beacon_round}


def derive_seed(fingerprint: str, payment: Payment) -> bytes:
    material = DOMAIN + lp(_hex32(fingerprint, "fingerprint")) + lp(_hex32(payment.txid, "txid"))
    for part in payment.entropy():
        material += lp(part)
    return hashlib.sha256(material).digest()


def check_canonical_order(payments: list) -> None:
    """Strictly increasing (height, position), unique txid.  No reordering is ever done silently."""
    seen = set()
    prev = None
    for p in payments:
        if p.txid in seen:
            raise DrawError(f"duplicate txid {p.txid}")
        seen.add(p.txid)
        key = (p.height, p.position)
        if prev is not None and key <= prev:
            raise DrawError(f"payments not in canonical order at {p.txid}: {key} after {prev}")
        prev = key


# --------------------------------------------------------------------------
# The draw
# --------------------------------------------------------------------------

@dataclass
class Result:
    fingerprint: str
    assignments: list = field(default_factory=list)   # one dict per NFT handed out
    refunds: list = field(default_factory=list)       # one dict per (partial) refund
    entropy: list = field(default_factory=list)       # one dict per drawn payment (seed inputs, archived)
    remaining: dict = field(default_factory=dict)

    def to_json(self) -> dict:
        body = {"pot_fingerprint": self.fingerprint, "assignments": self.assignments,
                "refunds": self.refunds, "entropy": self.entropy, "remaining": self.remaining}
        body["ledger_digest"] = ledger_digest(body)
        return body


def ledger_digest(body: dict) -> str:
    """SHA-256 over the canonical JSON of the ledger without its digest field."""
    core = {k: v for k, v in body.items() if k != "ledger_digest"}
    return hashlib.sha256(canonical_json(core)).hexdigest()


def run_draw(pot: PotList, payments: list, rules: Rules = Rules()) -> Result:
    check_canonical_order(payments)
    fp = pot.fingerprint()
    free = {GROUP_POT: list(pot.pot), GROUP_EPIC: list(pot.epics), GROUP_ULTRA: list(pot.ultras)}
    res = Result(fingerprint=fp)

    def refund(p: Payment, amount: int, reason: str) -> None:
        net = amount - rules.refund_fee
        res.refunds.append({"txid": p.txid, "height": p.height, "position": p.position, "sender": p.sender,
                            "paid": p.amount, "refund": max(net, 0), "kept": min(amount, rules.refund_fee),
                            "reason": reason})

    for p in payments:
        group, n = rules.classify(p.amount)
        if group is None:
            refund(p, p.amount, "amount_not_a_price")
            continue
        available = len(free[group])
        take = min(n, available) if rules.partial_fill else (n if available >= n else 0)
        if take == 0:
            refund(p, p.amount, f"{group}_empty")
            continue
        if rules.require_seed_block and not p.seed_block_hash:
            raise DrawError(f"payment {p.txid}: seed_block_hash (hash of block P+{SEED_OFFSET}) is required")
        if rules.require_beacon and not p.beacon:
            raise DrawError(f"payment {p.txid}: beacon value is required")
        seed = derive_seed(fp, p)
        res.entropy.append(p.entropy_record())
        for j in range(take):
            idx = rand_below(seed, group, j, len(free[group]))
            item = free[group].pop(idx)
            res.assignments.append({"txid": p.txid, "height": p.height, "position": p.position, "sender": p.sender,
                                    "group": group, "j": j, "nft_id": item.id, "rarity": item.rarity,
                                    "meta_hash": item.meta_hash})
        if take < n:  # partial fill (only reachable with partial_fill=True)
            refund(p, (n - take) * rules.unit_price, "pot_partially_empty")
    res.remaining = {g: len(v) for g, v in free.items()}
    return res


# --------------------------------------------------------------------------
# CLI
# --------------------------------------------------------------------------

def load_rules(path) -> Rules:
    if not path:
        return Rules()
    with open(path, encoding="utf-8") as fh:
        doc = json.load(fh)
    if not isinstance(doc, dict):
        raise DrawError("rules file must be a JSON object")
    unknown = set(doc) - set(Rules().to_json())
    if unknown:
        raise DrawError(f"rules: unknown keys {sorted(unknown)}")
    return Rules(**doc)


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description="PCAT 702 reference draw (offline, deterministic)")
    ap.add_argument("--pot", required=True, help="pot list JSON")
    ap.add_argument("--payments", required=True, help="payments JSON (canonical order, with seed_block_hash + beacon)")
    ap.add_argument("--rules", help="rules JSON (defaults: 23/223/523 NOID, entropy required)")
    ap.add_argument("--expect-pcat702", action="store_true", help="enforce 414/191/69/15/1 + 8 + 4")
    ap.add_argument("--expect-fingerprint", metavar="HEX", help="refuse any pot list with another fingerprint")
    ap.add_argument("--out", help="write result JSON here (default stdout)")
    args = ap.parse_args(argv)
    try:
        with open(args.pot, encoding="utf-8") as fh:
            pot = PotList.from_json(json.load(fh))
        if args.expect_pcat702:
            pot.check_pcat702()
        if args.expect_fingerprint and pot.fingerprint() != args.expect_fingerprint.lower():
            raise DrawError(f"pot fingerprint {pot.fingerprint()} != expected {args.expect_fingerprint}")
        with open(args.payments, encoding="utf-8") as fh:
            payments = [Payment.from_json(d) for d in json.load(fh)]
        result = run_draw(pot, payments, load_rules(args.rules)).to_json()
    except (DrawError, OSError, json.JSONDecodeError, TypeError, ValueError) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2
    text = json.dumps(result, indent=1, sort_keys=True)
    if args.out:
        with open(args.out, "w", encoding="utf-8") as fh:
            fh.write(text + "\n")
    else:
        print(text)
    return 0


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