#!/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())