order_reconciler.py

"""Order reconciliation for GTC PostOnly quotes.

The reconciler compares desired quote targets with live open orders and decides
which orders can keep their queue position, which stale orders should be
cancelled, and which missing targets should be posted.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from decimal import Decimal
from typing import Any

from config import BotConfig
from polymarket_adapter import PolymarketAdapter, TokenConfig
from quote_engine import TargetOrder


LOGGER = logging.getLogger(__name__)


@dataclass(frozen=True)
class ReconciliationPlan:
    keep: list[TargetOrder]
    place: list[TargetOrder]
    cancel_order_ids: list[str]
    cancel_all: bool = False
    reason: str = ""

    @property
    def has_work(self) -> bool:
        return bool(self.place or self.cancel_order_ids or self.cancel_all)


class OrderReconciler:
    def __init__(self, config: BotConfig, adapter: PolymarketAdapter) -> None:
        self.config = config
        self.adapter = adapter

    def plan(
        self,
        desired_orders: list[TargetOrder],
        last_targets: list[TargetOrder],
        tokens: tuple[TokenConfig, TokenConfig],
        *,
        condition_id: str,
        use_live_orders: bool,
    ) -> ReconciliationPlan:
        if not use_live_orders:
            return self._plan_from_last_targets(desired_orders, last_targets)

        try:
            open_orders = self._load_open_orders(tokens)
        except Exception as exc:
            LOGGER.warning("[RECONCILE] Could not load open orders; using last target fallback: %s", exc)
            return self._plan_from_last_targets(desired_orders, last_targets)

        desired_by_key = {self._target_key(order): order for order in desired_orders}
        live_by_key: dict[tuple[str, str], list[dict[str, Any]]] = {}
        stale_order_ids: list[str] = []

        for row in open_orders:
            key = self._open_order_key(row)
            if key is None:
                order_id = self._open_order_id(row)
                if order_id:
                    stale_order_ids.append(order_id)
                continue
            live_by_key.setdefault(key, []).append(row)

        keep: list[TargetOrder] = []
        place: list[TargetOrder] = []
        cancel_order_ids = list(stale_order_ids)

        for key, desired in desired_by_key.items():
            matching_row = self._find_matching_open_order(desired, live_by_key.get(key, []))
            if matching_row is None:
                place.append(desired)
                cancel_order_ids.extend(
                    order_id
                    for order_id in (self._open_order_id(row) for row in live_by_key.get(key, []))
                    if order_id
                )
            else:
                keep.append(desired)
                keep_id = self._open_order_id(matching_row)
                for row in live_by_key.get(key, []):
                    order_id = self._open_order_id(row)
                    if order_id and order_id != keep_id:
                        cancel_order_ids.append(order_id)

        for key, rows in live_by_key.items():
            if key in desired_by_key:
                continue
            cancel_order_ids.extend(
                order_id for order_id in (self._open_order_id(row) for row in rows) if order_id
            )

        return ReconciliationPlan(
            keep=keep,
            place=place,
            cancel_order_ids=sorted(set(cancel_order_ids)),
            reason=f"live keep={len(keep)} place={len(place)} cancel={len(set(cancel_order_ids))}",
        )

    def _plan_from_last_targets(
        self,
        desired_orders: list[TargetOrder],
        last_targets: list[TargetOrder],
    ) -> ReconciliationPlan:
        if self._same_targets(desired_orders, last_targets):
            return ReconciliationPlan(keep=desired_orders, place=[], cancel_order_ids=[], reason="unchanged")
        return ReconciliationPlan(
            keep=[],
            place=desired_orders,
            cancel_order_ids=[],
            cancel_all=bool(last_targets),
            reason="target_changed",
        )

    def _load_open_orders(self, tokens: tuple[TokenConfig, TokenConfig]) -> list[dict[str, Any]]:
        rows: list[dict[str, Any]] = []
        for token in tokens:
            rows.extend(self.adapter.open_orders_for_token(token.token_id))
        return rows

    @staticmethod
    def _same_targets(new_orders: list[TargetOrder], old_orders: list[TargetOrder]) -> bool:
        if len(new_orders) != len(old_orders):
            return False
        old_map = {OrderReconciler._target_key(order): order for order in old_orders}
        for new_order in new_orders:
            old_order = old_map.get(OrderReconciler._target_key(new_order))
            if old_order is None:
                return False
            if abs(new_order.price - old_order.price) >= new_order.token.tick_size:
                return False
            if abs(new_order.shares - old_order.shares) > Decimal("0.000001"):
                return False
        return True

    @staticmethod
    def _target_key(order: TargetOrder) -> tuple[str, str]:
        return order.token.token_id, order.side.upper()

    @staticmethod
    def _open_order_key(row: dict[str, Any]) -> tuple[str, str] | None:
        token_id = str(
            row.get("asset_id")
            or row.get("asset")
            or row.get("token_id")
            or row.get("tokenId")
            or ""
        )
        side = str(row.get("side") or "").upper()
        if not token_id or side not in {"BUY", "SELL"}:
            return None
        return token_id, side

    @staticmethod
    def _open_order_id(row: dict[str, Any]) -> str:
        return str(row.get("id") or row.get("order_id") or row.get("orderId") or "")

    @staticmethod
    def _open_order_price(row: dict[str, Any]) -> Decimal:
        return Decimal(str(row.get("price") or row.get("limit_price") or "0"))

    @staticmethod
    def _open_order_size(row: dict[str, Any]) -> Decimal:
        return Decimal(
            str(
                row.get("remaining_size")
                or row.get("remainingSize")
                or row.get("size")
                or row.get("amount")
                or row.get("original_size")
                or "0"
            )
        )

    def _find_matching_open_order(
        self,
        desired: TargetOrder,
        rows: list[dict[str, Any]],
    ) -> dict[str, Any] | None:
        """
        Tries to find a live order that is "close enough" to the desired target.
        If found, we can keep it to defend our queue position.
        """
        for row in rows:
            price = self._open_order_price(row)
            size = self._open_order_size(row)

            # 1. Price tolerance check
            price_tolerance = self.config.reconcile_price_tolerance_bps / Decimal("10000")
            # Avoid division by zero if desired price is 0
            if desired.price > 0 and (abs(price - desired.price) / desired.price > price_tolerance):
                continue

            # 2. Size tolerance check
            size_tolerance = self.config.reconcile_size_tolerance_pct / Decimal("100")
            # Avoid division by zero if desired shares is 0
            if desired.shares > 0 and (abs(size - desired.shares) / desired.shares > size_tolerance):
                continue

            # If both checks pass, we have found a suitable order to keep.
            LOGGER.info(
                "[RECONCILE_KEEP] Found close enough live order to keep for %s %s. Target: %s @ %s, Live: %s @ %s",
                desired.side,
                desired.token.outcome,
                desired.shares,
                desired.price,
                size,
                price,
            )
            return row

        # No suitable live order found to keep for this target.
        return None