#!/usr/bin/env python3
"""
WebSocket Trading Bot for PolySimulator

Streams real-time prices via WebSocket and places orders via REST
when price signals are detected.

Usage:
    export POLYSIM_API_KEY="ps_live_..."
    python ws_bot.py
"""

import asyncio
import json
import os
import time

import aiohttp


API_KEY = os.environ["POLYSIM_API_KEY"]
BASE_URL = os.environ.get("POLYSIM_BASE_URL", "https://api.polysimulator.com")
WS_URL = BASE_URL.replace("http://", "ws://").replace("https://", "wss://")


async def get_ws_token(session):
    """Obtain a short-lived WebSocket authentication token."""
    async with session.post(
        f"{BASE_URL}/v1/keys/ws-token",
        headers={"X-API-Key": API_KEY},
    ) as resp:
        data = await resp.json()
        return data["token"]


async def place_order(session, market_id, side, outcome, quantity, price):
    """Place an order via REST API.

    Note: ``price`` is REQUIRED even for market orders — it acts as a
    worst-price slippage cap. Pass it as a string. Compute from the
    live WS frame, and **always derive the cap from the price of the
    outcome you are trading**: for BUY Yes use ``min(0.99,
    float(data["buy"]) * 1.05)``, for SELL Yes use ``max(0.01,
    float(data["buy"]) * 0.95)``. ``data["buy"]`` is the Yes price;
    ``data["sell"]`` is the No price — *not* a "sell-side ask" for Yes.
    """
    payload = {
        "market_id": market_id,
        "side": side,
        "outcome": outcome,
        "quantity": str(quantity),
        "price": str(price),           # Worst-price slippage cap
        "order_type": "market",
    }
    async with session.post(
        f"{BASE_URL}/v1/orders",
        headers={
            "X-API-Key": API_KEY,
            "Content-Type": "application/json",
            "Idempotency-Key": f"ws-{market_id}-{side}-{int(time.time())}",
        },
        json=payload,
    ) as resp:
        return await resp.json()


async def run_ws_bot(market_ids: list[str]):
    """Main WebSocket bot loop."""
    async with aiohttp.ClientSession() as session:
        # 1. Get WebSocket token
        token = await get_ws_token(session)
        print(f"Got WS token, connecting...")

        # 2. Connect to WebSocket
        async with session.ws_connect(
            f"{WS_URL}/v1/ws/prices?token={token}"
        ) as ws:
            # 3. Subscribe to markets
            await ws.send_json({
                "action": "subscribe",
                "markets": market_ids,
            })
            print(f"Subscribed to {len(market_ids)} markets")

            # 4. Process price updates
            async for msg in ws:
                if msg.type == aiohttp.WSMsgType.TEXT:
                    data = json.loads(msg.data)

                    if data.get("type") == "price":
                        market_id = data["market_id"]
                        yes_price = float(data["buy"])

                        # Strategy: buy when Yes < 0.35, sell when > 0.65
                        if yes_price < 0.35:
                            # Worst-price BUY cap: 5% above live ask, clamped at 0.99
                            worst_buy = min(0.99, yes_price * 1.05)
                            print(f"BUY signal: {market_id} @ {yes_price:.2f} (cap {worst_buy:.2f})")
                            result = await place_order(
                                session, market_id, "BUY", "Yes", 5, worst_buy
                            )
                            print(f"  → {result.get('status')} @ {result.get('price')}")

                        elif yes_price > 0.65:
                            # Worst-price SELL floor: 5% below the live Yes
                            # quote, floored at 0.01. NOTE: derive the floor
                            # from `data["buy"]` (the Yes outcome price) — NOT
                            # `data["sell"]`, which is the No outcome price.
                            # A typical market with buy=0.70, sell=0.30 would
                            # otherwise produce a floor near 0.285, allowing
                            # fills far below the intended Yes valuation.
                            worst_sell = max(0.01, yes_price * 0.95)
                            print(f"SELL signal: {market_id} @ {yes_price:.2f} (floor {worst_sell:.2f})")
                            result = await place_order(
                                session, market_id, "SELL", "Yes", 5, worst_sell
                            )
                            print(f"  → {result.get('status')} @ {result.get('price')}")

                    elif data.get("type") == "pong":
                        pass  # Heartbeat response

                elif msg.type == aiohttp.WSMsgType.CLOSED:
                    print("WebSocket closed, reconnecting...")
                    break
                elif msg.type == aiohttp.WSMsgType.ERROR:
                    print(f"WebSocket error: {ws.exception()}")
                    break


async def main():
    # Define which markets to monitor
    market_ids = [
        "0x1234abcd...",  # Replace with actual condition IDs
        "0x5678efgh...",
    ]

    while True:
        try:
            await run_ws_bot(market_ids)
        except Exception as e:
            print(f"Connection error: {e}")
        print("Reconnecting in 5s...")
        await asyncio.sleep(5)


if __name__ == "__main__":
    asyncio.run(main())

const WebSocket = require("ws");
const fetch = require("node-fetch");

const API_KEY = process.env.POLYSIM_API_KEY;
const BASE_URL = process.env.POLYSIM_BASE_URL || "https://api.polysimulator.com";
const WS_URL = BASE_URL.replace("http://", "ws://").replace("https://", "wss://");

async function getWsToken() {
  const resp = await fetch(`${BASE_URL}/v1/keys/ws-token`, {
    method: "POST",
    headers: { "X-API-Key": API_KEY },
  });
  const data = await resp.json();
  return data.token;
}

// `price` is REQUIRED even for market orders — it acts as a worst-price
// slippage cap. For BUY Yes: `Math.min(0.99, data.buy * 1.05)`. For
// SELL Yes: `Math.max(0.01, data.buy * 0.95)` (the floor is on the
// SAME outcome being sold — `data.buy` is the Yes price; `data.sell`
// is the No price, not a "sell-side ask" for Yes). Pass it as a string.
async function placeOrder(marketId, side, outcome, quantity, price) {
  const resp = await fetch(`${BASE_URL}/v1/orders`, {
    method: "POST",
    headers: {
      "X-API-Key": API_KEY,
      "Content-Type": "application/json",
      "Idempotency-Key": `ws-${marketId}-${side}-${Date.now()}`,
    },
    body: JSON.stringify({
      market_id: marketId,
      side,
      outcome,
      quantity: String(quantity),
      price: String(price),     // Worst-price slippage cap
      order_type: "market",
    }),
  });
  return resp.json();
}

async function runBot(marketIds) {
  const token = await getWsToken();
  const ws = new WebSocket(`${WS_URL}/v1/ws/prices?token=${token}`);

  ws.on("open", () => {
    ws.send(JSON.stringify({
      action: "subscribe",
      markets: marketIds,
    }));
    console.log(`Subscribed to ${marketIds.length} markets`);
  });

  ws.on("message", async (raw) => {
    const data = JSON.parse(raw);
    if (data.type === "price") {
      const yesPrice = parseFloat(data.buy);
      if (yesPrice < 0.35) {
        const worstBuy = Math.min(0.99, yesPrice * 1.05);
        console.log(`BUY signal: ${data.market_id} @ ${yesPrice} (cap ${worstBuy.toFixed(2)})`);
        const result = await placeOrder(data.market_id, "BUY", "Yes", 5, worstBuy);
        console.log(`  → ${result.status} @ ${result.price}`);
      } else if (yesPrice > 0.65) {
        // Worst-price SELL floor: 5% below the live Yes quote, floored at
        // 0.01. NOTE: derive the floor from `data.buy` (the Yes outcome
        // price) — NOT `data.sell`, which is the No outcome price. A
        // typical market with buy=0.70, sell=0.30 would otherwise produce
        // a floor near 0.285, allowing fills far below the intended Yes
        // valuation.
        const worstSell = Math.max(0.01, yesPrice * 0.95);
        console.log(`SELL signal: ${data.market_id} @ ${yesPrice} (floor ${worstSell.toFixed(2)})`);
        const result = await placeOrder(data.market_id, "SELL", "Yes", 5, worstSell);
        console.log(`  → ${result.status} @ ${result.price}`);
      }
    }
  });

  ws.on("close", () => console.log("WebSocket closed"));
  ws.on("error", (err) => console.error("WebSocket error:", err));
}

runBot(["0x1234abcd...", "0x5678efgh..."]);