src/execution.py

import logging
from typing import Dict, List
from datetime import datetime, timezone
from src.strategies.arbitrage import ArbOpportunity

logger = logging.getLogger(__name__)

class PaperTradingEngine:
    def __init__(self, initial_capital: float = 10000.0):
        self.capital = initial_capital
        self.positions: Dict[str, Dict[str, float]] = {
            "polymarket": {},
            "kalshi": {}
        } # { platform: { market_id: size } }
        self.trade_history: List[dict] = []
        
        # Risk Limits
        self.max_position_size = 1000.0 # Max dollars per market
        self.daily_loss_limit = 500.0
        self.starting_capital_today = initial_capital

    def _check_risk_limits(self, platform: str, market_id: str, cost: float) -> bool:
        """Verify trade doesn't breach risk parameters."""
        current_exposure = self.positions[platform].get(market_id, 0.0) * cost # simplified exposure
        if cost > self.max_position_size:
            logger.warning(f"Risk Rejected: Trade size {cost} exceeds max {self.max_position_size}")
            return False
            
        daily_pnl = self.capital - self.starting_capital_today
        if daily_pnl < -self.daily_loss_limit:
             logger.warning(f"Risk Rejected: Daily loss limit {self.daily_loss_limit} breached.")
             return False
             
        return True

    def execute_arbitrage(self, opp: ArbOpportunity):
        """Execute a guaranteed arbitrage pair via paper trading."""
        # Calculate capital required
        buy_cost = opp.buy_price * opp.buy_size
        sell_margin_required = (1.0 - opp.sell_price) * opp.sell_size # Simplistic collateral assumption
        
        total_cost = buy_cost + sell_margin_required

        if self.capital < total_cost:
            logger.warning(f"Insufficient capital for Arb. Need {total_cost}, have {self.capital}")
            # Adjust size down to fit capital
            scale_factor = self.capital / total_cost
            opp.buy_size *= scale_factor
            opp.sell_size *= scale_factor
            
            # Recalculate
            buy_cost = opp.buy_price * opp.buy_size
            sell_margin_required = (1.0 - opp.sell_price) * opp.sell_size
            total_cost = buy_cost + sell_margin_required
            
            if opp.buy_size < 1: # Too small to trade
                return

        if not self._check_risk_limits(opp.buy_platform, opp.market_id_pm, buy_cost):
             return

        # Simulate Fills (Assume instant fill for now)
        # Deduct capital
        self.capital -= total_cost
        
        # Update Positions (Long buy side, short sell side)
        self.positions[opp.buy_platform][opp.market_id_pm] = self.positions[opp.buy_platform].get(opp.market_id_pm, 0) + opp.buy_size
        self.positions[opp.sell_platform][opp.market_id_kalshi] = self.positions[opp.sell_platform].get(opp.market_id_kalshi, 0) - opp.sell_size
        
        # Record Trade
        trade_record = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "strategy": "cross_platform_arb",
            "buy_leg": {"platform": opp.buy_platform, "market": opp.market_id_pm, "price": opp.buy_price, "size": opp.buy_size},
            "sell_leg": {"platform": opp.sell_platform, "market": opp.market_id_kalshi, "price": opp.sell_price, "size": opp.sell_size},
            "expected_profit": opp.expected_profit_margin * opp.buy_size,
            "capital_remaining": self.capital
        }
        self.trade_history.append(trade_record)
        
        logger.info(f"Paper Executed Arb! Expected Profit: ${trade_record['expected_profit']:.2f} | Capital: ${self.capital:.2f}")

    def get_portfolio_summary(self) -> dict:
        return {
            "capital_available": self.capital,
            "open_positions": self.positions,
            "total_trades": len(self.trade_history)
        }