polymarket_bot/execution.py

"""
Module d'exécution: placement d'ordres, gestion des positions.
Supporte le mode dry-run (paper trading) et le mode live.
"""
import asyncio
import logging
import time
from dataclasses import dataclass, field
from typing import Optional

from .config import BotConfig, POLYGON_CHAIN_ID, CLOB_API_URL

logger = logging.getLogger("polybot.execution")


# ══════════════════════════════════════════════════════════════════
# DATA MODELS
# ══════════════════════════════════════════════════════════════════
@dataclass
class Position:
    market_id: str
    token_id: str
    outcome: str          # "Yes" / "No"
    side: str             # "BUY"
    entry_price: float
    size: float           # Nombre de shares
    cost_basis: float     # Coût total en USD
    timestamp: float = 0.0
    strategy: str = ""
    pnl: float = 0.0

    @property
    def current_value(self) -> float:
        return self.size * self.entry_price

    def unrealized_pnl(self, current_price: float) -> float:
        return self.size * (current_price - self.entry_price)


@dataclass
class Trade:
    trade_id: str
    market_id: str
    token_id: str
    outcome: str
    side: str
    price: float
    size: float
    cost: float
    timestamp: float
    strategy: str
    status: str = "filled"  # filled, cancelled, rejected


@dataclass
class PortfolioState:
    balance_usd: float = 10000.0        # Capital initial simulé
    positions: dict = field(default_factory=dict)  # token_id -> Position
    trades: list = field(default_factory=list)
    total_pnl: float = 0.0
    daily_pnl: float = 0.0
    daily_pnl_reset_time: float = 0.0
    peak_balance: float = 10000.0
    max_drawdown: float = 0.0

    @property
    def total_exposure(self) -> float:
        return sum(p.cost_basis for p in self.positions.values())

    @property
    def available_capital(self) -> float:
        return self.balance_usd - self.total_exposure

    @property
    def num_positions(self) -> int:
        return len(self.positions)

    def update_drawdown(self):
        current_total = self.balance_usd + sum(
            p.unrealized_pnl(p.entry_price) for p in self.positions.values()
        )
        if current_total > self.peak_balance:
            self.peak_balance = current_total
        dd = (self.peak_balance - current_total) / self.peak_balance
        if dd > self.max_drawdown:
            self.max_drawdown = dd


# ══════════════════════════════════════════════════════════════════
# EXECUTION ENGINE
# ══════════════════════════════════════════════════════════════════
class ExecutionEngine:
    """
    Moteur d'exécution des ordres.
    Supporte dry-run (simulation) et live (via py-clob-client).
    """

    def __init__(self, config: BotConfig):
        self.config = config
        self.portfolio = PortfolioState()
        self._clob_client = None
        self._trade_counter = 0

    def _init_clob_client(self):
        """Initialise le client CLOB pour le mode live."""
        if self.config.dry_run:
            return

        try:
            from py_clob_client.client import ClobClient
            from py_clob_client.clob_types import ApiCreds

            # Step 1: Create client with wallet
            client = ClobClient(
                CLOB_API_URL,
                key=self.config.private_key,
                chain_id=POLYGON_CHAIN_ID,
            )

            # Step 2: Derive API credentials
            api_key_data = client.create_or_derive_api_key()
            creds = ApiCreds(
                api_key=api_key_data["apiKey"],
                api_secret=api_key_data["secret"],
                api_passphrase=api_key_data["passphrase"],
            )

            # Step 3: Reinit client with full auth
            self._clob_client = ClobClient(
                CLOB_API_URL,
                key=self.config.private_key,
                chain_id=POLYGON_CHAIN_ID,
                creds=creds,
            )
            logger.info("CLOB client initialized successfully (LIVE mode)")

        except Exception as e:
            logger.error(f"Failed to init CLOB client: {e}")
            logger.warning("Falling back to dry-run mode")
            self.config.dry_run = True

    # ── Risk Checks ──────────────────────────────────────────────
    def _check_risk(self, cost_usd: float, market_id: str) -> tuple[bool, str]:
        """Vérifie les contraintes de risque avant un trade."""

        # Check: capital disponible
        if cost_usd > self.portfolio.available_capital:
            return False, f"Insufficient capital: need ${cost_usd:.2f}, have ${self.portfolio.available_capital:.2f}"

        # Check: exposition totale
        if self.portfolio.total_exposure + cost_usd > self.config.max_total_exposure_usd:
            return False, f"Max total exposure reached: ${self.config.max_total_exposure_usd:.2f}"

        # Check: exposition par marché
        market_exposure = sum(
            p.cost_basis for p in self.portfolio.positions.values()
            if p.market_id == market_id
        )
        max_market = self.config.max_total_exposure_usd * self.config.max_single_market_exposure_pct
        if market_exposure + cost_usd > max_market:
            return False, f"Max market exposure reached: ${max_market:.2f}"

        # Check: nombre de positions
        if self.portfolio.num_positions >= self.config.max_concurrent_positions:
            return False, f"Max concurrent positions reached: {self.config.max_concurrent_positions}"

        # Check: perte journalière
        if self.portfolio.daily_pnl < -self.config.max_daily_loss_usd:
            return False, f"Daily loss limit reached: ${self.config.max_daily_loss_usd:.2f}"

        # Check: taille min/max
        if cost_usd < self.config.arb_min_position_usd:
            return False, f"Trade too small: ${cost_usd:.2f} < ${self.config.arb_min_position_usd:.2f}"

        return True, "OK"

    # ── Order Placement ──────────────────────────────────────────
    async def place_order(
        self,
        token_id: str,
        market_id: str,
        outcome: str,
        side: str,
        price: float,
        size: float,
        strategy: str,
        order_type: str = "GTC",
    ) -> Optional[Trade]:
        """
        Place un ordre. Retourne un Trade si réussi.

        Args:
            token_id: ID du token (YES ou NO)
            market_id: ID du marché
            outcome: "Yes" ou "No"
            side: "BUY" ou "SELL"
            price: Prix par share
            size: Nombre de shares
            strategy: Nom de la stratégie ("arbitrage", "value_bet", etc.)
            order_type: "GTC" (Good Till Cancel), "FOK" (Fill or Kill)
        """
        cost = price * size

        # Risk check
        can_trade, reason = self._check_risk(cost, market_id)
        if not can_trade:
            logger.warning(f"Trade rejected by risk manager: {reason}")
            return None

        self._trade_counter += 1
        trade_id = f"T{self._trade_counter:06d}"

        if self.config.dry_run:
            # ── DRY RUN: simulate fill ───────────────────────────
            trade = Trade(
                trade_id=trade_id,
                market_id=market_id,
                token_id=token_id,
                outcome=outcome,
                side=side,
                price=price,
                size=size,
                cost=cost,
                timestamp=time.time(),
                strategy=strategy,
                status="filled",
            )

            # Update portfolio
            self.portfolio.balance_usd -= cost
            self.portfolio.positions[token_id] = Position(
                market_id=market_id,
                token_id=token_id,
                outcome=outcome,
                side=side,
                entry_price=price,
                size=size,
                cost_basis=cost,
                timestamp=time.time(),
                strategy=strategy,
            )
            self.portfolio.trades.append(trade)

            logger.info(
                f"[DRY RUN] {side} {size:.1f} {outcome} @ ${price:.4f} "
                f"= ${cost:.2f} | Strategy: {strategy} | Market: {market_id[:16]}..."
            )
            return trade

        else:
            # ── LIVE MODE: place via CLOB ────────────────────────
            try:
                from py_clob_client.clob_types import OrderArgs, OrderType, BUY, SELL

                if self._clob_client is None:
                    self._init_clob_client()

                side_enum = BUY if side == "BUY" else SELL
                otype = OrderType.GTC if order_type == "GTC" else OrderType.FOK

                order = self._clob_client.create_order(OrderArgs(
                    token_id=token_id,
                    price=price,
                    size=size,
                    side=side_enum,
                    order_type=otype,
                ))
                resp = self._clob_client.post_order(order)

                if resp and resp.get("success"):
                    trade = Trade(
                        trade_id=resp.get("orderID", trade_id),
                        market_id=market_id,
                        token_id=token_id,
                        outcome=outcome,
                        side=side,
                        price=price,
                        size=size,
                        cost=cost,
                        timestamp=time.time(),
                        strategy=strategy,
                        status="filled",
                    )
                    self.portfolio.balance_usd -= cost
                    self.portfolio.positions[token_id] = Position(
                        market_id=market_id,
                        token_id=token_id,
                        outcome=outcome,
                        side=side,
                        entry_price=price,
                        size=size,
                        cost_basis=cost,
                        timestamp=time.time(),
                        strategy=strategy,
                    )
                    self.portfolio.trades.append(trade)

                    logger.info(
                        f"[LIVE] {side} {size:.1f} {outcome} @ ${price:.4f} "
                        f"= ${cost:.2f} | Strategy: {strategy}"
                    )
                    return trade
                else:
                    logger.error(f"Order placement failed: {resp}")
                    return None

            except Exception as e:
                logger.error(f"Order execution error: {e}")
                return None

    async def place_arb_pair(
        self,
        market_id: str,
        yes_token_id: str,
        no_token_id: str,
        yes_price: float,
        no_price: float,
        size: float,
    ) -> tuple[Optional[Trade], Optional[Trade]]:
        """
        Place une paire d'ordres d'arbitrage (BUY YES + BUY NO).
        Exécute simultanément pour minimiser le risque d'exécution partielle.
        """
        total_cost = (yes_price + no_price) * size
        can_trade, reason = self._check_risk(total_cost, market_id)
        if not can_trade:
            logger.warning(f"Arb pair rejected: {reason}")
            return None, None

        # Exécution simultanée
        yes_trade, no_trade = await asyncio.gather(
            self.place_order(
                yes_token_id, market_id, "Yes", "BUY",
                yes_price, size, "arbitrage", "FOK"
            ),
            self.place_order(
                no_token_id, market_id, "No", "BUY",
                no_price, size, "arbitrage", "FOK"
            ),
        )

        if yes_trade and no_trade:
            profit_per_share = 1.0 - yes_price - no_price
            total_profit = profit_per_share * size
            logger.info(
                f"✅ ARB FILLED: {size:.1f} shares | "
                f"YES@{yes_price:.4f} + NO@{no_price:.4f} = {yes_price+no_price:.4f} | "
                f"Expected profit: ${total_profit:.2f} ({profit_per_share*100:.1f}%)"
            )
        elif yes_trade or no_trade:
            logger.warning("⚠️ PARTIAL ARB FILL - one leg failed!")

        return yes_trade, no_trade

    # ── Position Management ──────────────────────────────────────
    def close_position(self, token_id: str, exit_price: float) -> Optional[float]:
        """Ferme une position et calcule le PnL."""
        if token_id not in self.portfolio.positions:
            return None

        pos = self.portfolio.positions[token_id]
        pnl = pos.size * (exit_price - pos.entry_price)
        pos.pnl = pnl
        self.portfolio.total_pnl += pnl
        self.portfolio.daily_pnl += pnl
        self.portfolio.balance_usd += pos.size * exit_price

        logger.info(
            f"Position closed: {pos.outcome} | "
            f"Entry: ${pos.entry_price:.4f} → Exit: ${exit_price:.4f} | "
            f"PnL: ${pnl:+.2f}"
        )

        del self.portfolio.positions[token_id]
        self.portfolio.update_drawdown()
        return pnl

    def get_portfolio_summary(self) -> dict:
        """Résumé du portefeuille."""
        return {
            "balance_usd": round(self.portfolio.balance_usd, 2),
            "total_exposure": round(self.portfolio.total_exposure, 2),
            "available_capital": round(self.portfolio.available_capital, 2),
            "num_positions": self.portfolio.num_positions,
            "total_pnl": round(self.portfolio.total_pnl, 2),
            "daily_pnl": round(self.portfolio.daily_pnl, 2),
            "max_drawdown": f"{self.portfolio.max_drawdown*100:.2f}%",
            "num_trades": len(self.portfolio.trades),
            "win_rate": self._calculate_win_rate(),
        }

    def _calculate_win_rate(self) -> str:
        closed_trades = [t for t in self.portfolio.trades if t.status == "filled"]
        if not closed_trades:
            return "N/A"
        wins = sum(1 for t in closed_trades if t.strategy == "arbitrage")
        total = len(closed_trades)
        return f"{wins/total*100:.1f}%" if total > 0 else "N/A"