#!/usr/bin/env python3
"""Read one Polymarket market and one order book, plus a separate Vultax research table.

Python 3.10+, standard library only. No credentials, account access or orders.

  python3 polymarket_snapshot.py --list 10            find a market slug (1 request)
  python3 polymarket_snapshot.py --slug MARKET_SLUG   save a snapshot (3 requests)
"""
import argparse
import csv
import io
import json
import math
import sys
from datetime import datetime, timezone
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.parse import urlencode, urlparse
from urllib.request import Request, urlopen

MAX_BYTES = 4 * 1024 * 1024
GAMMA = "https://gamma-api.polymarket.com"
CLOB = "https://clob.polymarket.com"
STUDY = "https://vultax.com/research/best-polymarket-traders-leaderboards"
DATASET_ID = "section-activity-is-the-first-filter-chart"
FEE_CALCULATOR = "https://vultax.com/tools/polymarket-fee-calculator"


class ExampleError(Exception):
    """A problem the reader can act on. Printed as one line, without a traceback."""


def now():
    return datetime.now(timezone.utc).isoformat()


def fetch_json(url):
    request = Request(url, headers={"User-Agent": "Vultax-Public-Research-Example/1.1", "Accept": "application/json"})
    try:
        with urlopen(request, timeout=20) as response:
            body = response.read(MAX_BYTES + 1)
            if len(body) > MAX_BYTES:
                raise ExampleError(f"The response from {url} exceeds this example's 4 MiB limit")
            return json.loads(body), {"url": url, "retrievedAt": now(), "httpStatus": response.status}
    except HTTPError as error:
        retry_after = error.headers.get("Retry-After")
        hint = f"; wait {retry_after} seconds before trying again" if retry_after else ""
        raise ExampleError(f"HTTP {error.code} from {url}{hint}. This example never retries on its own.") from error
    except (URLError, TimeoutError) as error:
        raise ExampleError(f"Could not reach {urlparse(url).netloc}: {getattr(error, 'reason', error)}. Check the connection and run it again.") from error
    except json.JSONDecodeError as error:
        raise ExampleError(f"{url} did not return JSON") from error


def array(value, field):
    parsed = json.loads(value) if isinstance(value, str) else value
    if not isinstance(parsed, list):
        raise ValueError(f"{field} must be an array or a JSON-encoded array")
    return parsed


def finite_number(value):
    if value is None or isinstance(value, bool) or value == "":
        return None
    try:
        result = float(value)
    except (TypeError, ValueError):
        return None
    return result if math.isfinite(result) else None


def slug_from(value):
    """Accepts a slug or a pasted polymarket.com address and returns its last path segment."""
    text = value.strip()
    if "://" in text or text.startswith("polymarket.com"):
        path = urlparse(text if "://" in text else "https://" + text).path
        text = path.rstrip("/").rsplit("/", 1)[-1]
    if not text or any(character.isspace() for character in text):
        raise ExampleError("Pass a market slug such as will-x-happen-by-2026, or the market's polymarket.com address")
    return text


def normalize_market(market):
    labels = array(market.get("outcomes"), "outcomes")
    token_ids = array(market.get("clobTokenIds"), "clobTokenIds")
    raw_prices = market.get("outcomePrices")
    prices = [None] * len(labels) if raw_prices is None else array(raw_prices, "outcomePrices")
    if not labels or len(labels) != len(token_ids) or len(labels) != len(prices):
        raise ValueError("Outcome, price and token arrays do not align; refusing to guess their mapping")
    if any(not isinstance(token, str) or not token.isdigit() for token in token_ids):
        raise ValueError("Token IDs must remain decimal strings; numeric IDs can lose precision")
    outcomes = []
    for label, token, raw_price in zip(labels, token_ids, prices):
        price = finite_number(raw_price)
        if price is not None and not 0 <= price <= 1:
            raise ValueError("Outcome price is outside [0, 1]")
        outcomes.append({"label": str(label), "tokenId": token, "gammaPrice": price})
    return {
        "id": str(market["id"]), "conditionId": market.get("conditionId"), "slug": market.get("slug"),
        "question": market.get("question"), "sourceUpdatedAt": market.get("updatedAt"),
        "active": market.get("active"), "closed": market.get("closed"),
        "acceptingOrders": market.get("acceptingOrders"), "outcomes": outcomes,
        "feesEnabled": market.get("feesEnabled"), "feeType": market.get("feeType"), "feeSchedule": market.get("feeSchedule"),
    }


def choose_outcome(outcomes, wanted):
    """The first outcome by default, under its real label. Index 0 is never assumed to mean Yes."""
    if wanted is None:
        return outcomes[0]
    matches = [row for row in outcomes if row["label"].casefold() == wanted.casefold()]
    if len(matches) != 1:
        labels = ", ".join(repr(row["label"]) for row in outcomes)
        raise ExampleError(f"--outcome {wanted!r} does not name exactly one outcome of this market. Its outcomes are: {labels}")
    return matches[0]


def normalize_book(book, token_id, condition_id):
    if str(book.get("asset_id")) != token_id or book.get("market") != condition_id:
        raise ValueError("Book identity does not match the requested outcome and condition")
    def levels(side):
        result = []
        for level in book.get(side, []):
            price, size = finite_number(level.get("price")), finite_number(level.get("size"))
            if price is None or size is None or not 0 < price < 1 or size <= 0:
                raise ValueError(f"Invalid {side} level; refusing to silently discard it")
            result.append({"price": price, "size": size})
        return result
    bids, asks = levels("bids"), levels("asks")
    best_bid = max((row["price"] for row in bids), default=None)
    best_ask = min((row["price"] for row in asks), default=None)
    spread = round(best_ask - best_bid, 6) if best_bid is not None and best_ask is not None else None
    stamp = finite_number(book.get("timestamp"))
    return {"tokenId": token_id, "conditionId": condition_id, "sourceTimestampMs": book.get("timestamp"),
            "sourceTimestamp": datetime.fromtimestamp(stamp / 1000, timezone.utc).isoformat() if stamp else None,
            "bestBid": best_bid, "bestAsk": best_ask, "spread": spread,
            "crossedBook": spread < 0 if spread is not None else None,
            "bidLevels": len(bids), "askLevels": len(asks),
            "bestBidSize": sum(row["size"] for row in bids if row["price"] == best_bid) if bids else None,
            "bestAskSize": sum(row["size"] for row in asks if row["price"] == best_ask) if asks else None}


def taker_fee(shares, price, rate):
    """Polymarket's published formula, rounded once to five decimal places."""
    return round(shares * rate * price * (1 - price) + 1e-12, 5)


def fee_estimate(market, book, shares=100):
    """A taker fee for the best ask, from the market's own fee fields. Anything unknown stays unknown."""
    schedule = market.get("feeSchedule") if isinstance(market.get("feeSchedule"), dict) else None
    rate = finite_number(schedule.get("rate")) if schedule else None
    if market.get("feesEnabled") is not True:
        return {"shares": shares, "takerFee": 0.0, "basis": "feesEnabled is not true on this market", "calculator": FEE_CALCULATOR}
    if rate == 0:
        return {"shares": shares, "takerFee": 0.0, "basis": "this market's rate is 0", "calculator": FEE_CALCULATOR}
    if rate is None or schedule.get("exponent") != 1:
        return {"shares": shares, "takerFee": None, "basis": "fees are enabled but the rate is missing or the exponent is not 1; not estimated"}
    if book is None or book["bestAsk"] is None:
        return {"shares": shares, "rate": rate, "takerFee": None, "basis": "no ask to price the example at"}
    return {"shares": shares, "price": book["bestAsk"], "rate": rate, "takerFee": taker_fee(shares, book["bestAsk"], rate),
            "basis": "feeSchedule.rate on this market, exponent 1, at the best ask; an estimate, not a quote",
            "calculator": f"{FEE_CALCULATOR}?rate={rate:g}&shares={shares}&price={book['bestAsk'] * 100:g}"}


def select_research(document):
    if document.get("articleUrl") != STUDY:
        raise ValueError("Unexpected research article identity")
    dataset = next((row for row in document.get("datasets", []) if row.get("id") == DATASET_ID), None)
    if dataset is None or dataset.get("isLive") is not False:
        raise ValueError("Expected published snapshot dataset is absent")
    if not dataset.get("version") or not dataset.get("measurement"):
        raise ValueError("Research version or measurement context is missing")
    if any(len(row) != len(dataset["columns"]) for row in dataset["rows"]):
        raise ValueError("Research columns and rows do not align")
    return dataset


def csv_cell(value):
    text = "" if value is None else str(value)
    if isinstance(value, str) and text.lstrip().startswith(("=", "+", "-", "@")):
        return "'" + text
    return text


def render_csv(columns, rows):
    output = io.StringIO(newline="")
    writer = csv.writer(output)
    writer.writerow(columns)
    writer.writerows([[csv_cell(cell) for cell in row] for row in rows])
    return output.getvalue()


def list_markets(count):
    url = f"{GAMMA}/markets?" + urlencode({"limit": count, "active": "true", "closed": "false", "order": "volume24hr", "ascending": "false"})
    markets, _receipt = fetch_json(url)
    if not isinstance(markets, list):
        raise ExampleError("Gamma did not return a list of markets")
    print(f"{len(markets)} open markets by 24-hour volume, read {now()}\n")
    for market in markets:
        schedule = market.get("feeSchedule") if isinstance(market.get("feeSchedule"), dict) else {}
        fee = "no fee" if market.get("feesEnabled") is not True or schedule.get("rate") == 0 else f"rate {schedule.get('rate')}"
        print(f"{market.get('slug')}\n    {market.get('question')}  [{fee}]")
    print("\nNext: python3 polymarket_snapshot.py --slug ONE_OF_THE_SLUGS_ABOVE")


def find_market(slug):
    markets, receipt = fetch_json(f"{GAMMA}/markets?" + urlencode({"slug": slug}))
    if isinstance(markets, list) and len(markets) == 1 and markets[0].get("slug") == slug:
        return markets[0], receipt
    # Not a market. If it names an event, show that event's markets and stop; never pick one for the reader.
    events, _ = fetch_json(f"{GAMMA}/events?" + urlencode({"slug": slug}))
    children = [m.get("slug") for m in (events[0].get("markets") or [])] if isinstance(events, list) and len(events) == 1 else []
    if children:
        shown = "\n  ".join(children[:20])
        raise ExampleError(f"{slug!r} is an event with {len(children)} markets, not one market. Pass one of its market slugs:\n  {shown}")
    raise ExampleError(f"No market has the exact slug {slug!r}. Run with --list 10 to see current slugs.")


def snapshot(slug, wanted_outcome, output):
    raw_market, gamma_receipt = find_market(slug)
    market = normalize_market(raw_market)
    outcome = choose_outcome(market["outcomes"], wanted_outcome)
    book = None
    book_receipt = {"status": "unavailable", "reason": "market is not confirmed open for orders"}
    if market["active"] is True and market["closed"] is False and market["acceptingOrders"] is True:
        book_url = f"{CLOB}/book?" + urlencode({"token_id": outcome["tokenId"]})
        try:
            raw_book, book_receipt = fetch_json(book_url)
            book = normalize_book(raw_book, outcome["tokenId"], market["conditionId"])
        except (ExampleError, ValueError) as error:
            book_receipt = {"url": book_url, "retrievedAt": now(), "status": "unavailable", "reason": str(error)}
    research, research_receipt = fetch_json(STUDY + "/data.json")
    dataset = select_research(research)
    result = {"schemaVersion": "vultax-public-api-example-v2", "generatedAt": now(),
              "market": market, "bookOutcome": outcome["label"], "book": book,
              "feeEstimate": fee_estimate(market, book), "researchDataset": dataset,
              "receipts": {"gamma": gamma_receipt, "book": book_receipt, "research": research_receipt},
              "limitations": "Sequential public snapshots, not an atomic quote or executable trade. Gamma price is separate from CLOB bid/ask. Missing values remain null. The fee is an estimate from the market's reported settings. Research is historical and is not joined to this market."}
    output.mkdir(parents=True, exist_ok=False)
    (output / "snapshot.json").write_text(json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False) + "\n", encoding="utf-8")
    (output / "outcomes.csv").write_text(render_csv(["market_id", "condition_id", "outcome", "token_id", "gamma_price", "source_updated_at"],
        [[market["id"], market["conditionId"], row["label"], row["tokenId"], row["gammaPrice"], market["sourceUpdatedAt"]] for row in market["outcomes"]]), encoding="utf-8", newline="")
    (output / "research.csv").write_text(render_csv([column["name"] for column in dataset["columns"]], dataset["rows"]), encoding="utf-8", newline="")
    return result


def describe(result, output):
    market, book, fee = result["market"], result["book"], result["feeEstimate"]
    show = lambda value: "none" if value is None else f"{value:g}"
    print(market["question"])
    for row in market["outcomes"]:
        print(f"  {row['label']:<24} Gamma price {show(row['gammaPrice'])}")
    if book is None:
        print(f"Order book for {result['bookOutcome']!r}: unavailable ({result['receipts']['book'].get('reason')})")
    else:
        bid = f"bid {book['bestBid']:g} ({book['bestBidSize']:,.0f} shares)" if book["bestBid"] is not None else "no bids"
        ask = f"ask {book['bestAsk']:g} ({book['bestAskSize']:,.0f} shares)" if book["bestAsk"] is not None else "no asks"
        spread = f"spread {book['spread']:g}" if book["spread"] is not None else "no spread without both sides"
        print(f"Order book for {result['bookOutcome']!r}: {bid}, {ask}, {spread}; book time {book['sourceTimestamp']}")
    if fee["takerFee"] is None:
        print(f"Taker fee for {fee['shares']} shares: not estimated ({fee['basis']})")
    elif "rate" in fee:
        print(f"Taker fee for {fee['shares']} shares at the {fee['price']:g} ask: ${fee['takerFee']:g} (this market's rate is {fee['rate']:g}; an estimate, not a quote)")
    else:
        print(f"Taker fee for {fee['shares']} shares: $0 ({fee['basis']})")
    print(f"Research table: {len(result['researchDataset']['rows'])} rows, {result['researchDataset']['measurement'].get('population')}, kept separate from this market")
    print(f"Saved snapshot.json, outcomes.csv and research.csv to {output}")


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument("--slug", help="Exact market slug, or the market's polymarket.com address")
    mode.add_argument("--list", type=int, metavar="N", help="Print the N most-traded open markets (1 to 100) with their slugs, then stop")
    parser.add_argument("--outcome", help="Outcome label whose book to read, for example No (default: the first outcome)")
    parser.add_argument("--output", type=Path, help="New output directory; an existing one is never overwritten")
    args = parser.parse_args()
    try:
        if args.list is not None:
            if not 1 <= args.list <= 100:
                parser.error("--list takes a number from 1 to 100")
            return list_markets(args.list)
        output = args.output or Path("polymarket-example-" + datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ"))
        if output.exists():
            parser.error("Output directory already exists; choose a new path")
        describe(snapshot(slug_from(args.slug), args.outcome, output), output)
    except (ExampleError, ValueError) as error:
        sys.exit(f"error: {error}")


if __name__ == "__main__":
    main()
