execution.py

"""
execution.py
------------
Handles all order execution against the Hyperliquid REST API.

Responsibilities:
- Place market orders (entry).
- Place limit orders (TP/SL as reduce-only orders).
- Cancel open orders.
- Query open positions and account equity.
- Wrap Hyperliquid SDK calls with error handling and logging.
"""

import logging
import time
from dataclasses import dataclass
from typing import Optional

from config import Config
from strategy_flip import Direction
from responses import (
    FillResponse, CancelResponse, AccountState,
    ResponseStatus, TradeResponse,
)

logger = logging.getLogger(__name__)

# Hyperliquid Python SDK imports
try:
    from hyperliquid.exchange import Exchange
    from hyperliquid.info import Info
    from hyperliquid.utils import constants
    import eth_account
    HL_SDK_AVAILABLE = True
except ImportError:
    logger.warning("[Execution] hyperliquid-python-sdk not installed. Running in mock mode.")
    HL_SDK_AVAILABLE = False


@dataclass
class OrderResult:
    """
    Legacy shim — kept for backward compatibility with controller.py.
    New code should use FillResponse from responses.py directly.
    """
    success: bool
    order_id: Optional[str]
    filled_price: Optional[float]
    filled_size: Optional[float]
    error: Optional[str] = None
    raw_response: Optional[dict] = None
    fill: Optional[FillResponse] = None   # ← the typed response object

    @classmethod
    def from_fill(cls, fill: FillResponse) -> "OrderResult":
        """Convert a FillResponse into an OrderResult."""
        return cls(
            success=fill.is_ok() or fill.status == ResponseStatus.DRY_RUN,
            order_id=fill.order_id,
            filled_price=fill.filled_price,
            filled_size=fill.filled_size,
            error=fill.error,
            raw_response=fill.raw,
            fill=fill,
        )


@dataclass
class Position:
    symbol: str
    direction: Direction
    size: float
    entry_price: float
    unrealized_pnl: float
    leverage: float


class ExecutionEngine:
    """
    Wraps Hyperliquid SDK for order placement and account queries.
    Falls back to mock/dry-run mode if SDK is unavailable or DRY_RUN=true.
    """

    def __init__(self, config: Config, data_feed=None):
        self.config = config
        self.dry_run = config.DRY_RUN
        self._exchange: Optional[object] = None
        self._info: Optional[object] = None
        self._wallet_address = config.HL_WALLET_ADDRESS
        self._data_feed = data_feed  # used for live price in dry-run mode

        if not self.dry_run and HL_SDK_AVAILABLE:
            self._init_sdk()
        else:
            mode = "DRY RUN" if self.dry_run else "MOCK (SDK missing)"
            logger.warning(f"[Execution] Running in {mode} mode — no real orders will be placed.")

    # ------------------------------------------------------------------
    # SDK Initialization
    # ------------------------------------------------------------------

    def _init_sdk(self) -> None:
        """Initialize Hyperliquid SDK with wallet credentials."""
        try:
            account = eth_account.Account.from_key(self.config.HL_API_SECRET)
            self._info = Info(self.config.HL_REST_URL, skip_ws=True)
            self._exchange = Exchange(
                account,
                self.config.HL_REST_URL,
                account_address=self.config.HL_WALLET_ADDRESS,
            )
            logger.info(f"[Execution] Hyperliquid SDK initialized. Wallet: {self.config.HL_WALLET_ADDRESS[:8]}...")
        except Exception as e:
            logger.error(f"[Execution] SDK initialization failed: {e}")
            self._exchange = None
            self._info = None

    # ------------------------------------------------------------------
    # Account Queries
    # ------------------------------------------------------------------

    def get_account_equity(self) -> float:
        """Return current account equity in USD."""
        if self.dry_run or not self._info:
            logger.debug("[Execution] Mock equity: $10,000.00")
            return 10_000.0

        try:
            state = self._info.user_state(self._wallet_address)
            equity = float(state.get("marginSummary", {}).get("accountValue", 0))
            logger.debug(f"[Execution] Account equity: ${equity:.2f}")
            return equity
        except Exception as e:
            logger.error(f"[Execution] Failed to fetch account equity: {e}")
            return 0.0

    def get_open_positions(self) -> list[Position]:
        """Return list of currently open positions."""
        if self.dry_run or not self._info:
            return []

        try:
            state = self._info.user_state(self._wallet_address)
            positions = []
            for pos in state.get("assetPositions", []):
                p = pos.get("position", {})
                size = float(p.get("szi", 0))
                if size == 0:
                    continue
                direction = Direction.LONG if size > 0 else Direction.SHORT
                positions.append(Position(
                    symbol=p.get("coin", ""),
                    direction=direction,
                    size=abs(size),
                    entry_price=float(p.get("entryPx", 0)),
                    unrealized_pnl=float(p.get("unrealizedPnl", 0)),
                    leverage=float(p.get("leverage", {}).get("value", 1)),
                ))
            return positions
        except Exception as e:
            logger.error(f"[Execution] Failed to fetch positions: {e}")
            return []

    def get_open_position_count(self) -> int:
        return len(self.get_open_positions())

    # ------------------------------------------------------------------
    # Order Placement
    # ------------------------------------------------------------------

    def set_leverage(self, leverage: int) -> bool:
        """Set leverage for the trading symbol."""
        if self.dry_run or not self._exchange:
            logger.info(f"[Execution] [DRY RUN] Set leverage {leverage}x for {self.config.SYMBOL}")
            return True
        try:
            result = self._exchange.update_leverage(leverage, self.config.SYMBOL, is_cross=True)
            logger.info(f"[Execution] Leverage set to {leverage}x for {self.config.SYMBOL}")
            return True
        except Exception as e:
            logger.error(f"[Execution] Failed to set leverage: {e}")
            return False

    def place_market_order(
        self,
        direction: Direction,
        size_contracts: float,
    ) -> OrderResult:
        """
        Place a market order.

        Args:
            direction: LONG or SHORT.
            size_contracts: Number of contracts to trade.

        Returns:
            OrderResult wrapping a FillResponse.
            Access the full typed response via result.fill
        """
        is_buy = direction == Direction.LONG
        action_str = "BUY" if is_buy else "SELL"

        if self.dry_run or not self._exchange:
            mock_price = 0.0
            if self._data_feed is not None:
                mock_price = self._data_feed.get_latest_price() or 0.0
            if mock_price <= 0:
                logger.warning("[Execution] [DRY RUN] Could not get live price — using 0.0 as fill price.")

            fill = FillResponse.dry_run(
                order_type="market_entry",
                direction=direction.value,
                symbol=self.config.SYMBOL,
                mock_price=mock_price,
                size=size_contracts,
            )
            logger.info(f"[Execution] {fill.summary()}")
            return OrderResult.from_fill(fill)

        try:
            result = self._exchange.market_open(
                self.config.SYMBOL,
                is_buy,
                size_contracts,
                slippage=0.001,
            )
            logger.debug(f"[Execution] Raw market order response: {result}")

            fill = FillResponse.from_hl_response(
                raw=result,
                order_type="market_entry",
                direction=direction.value,
                symbol=self.config.SYMBOL,
            )
            logger.info(f"[Execution] {fill.summary()}")
            return OrderResult.from_fill(fill)

        except Exception as e:
            logger.error(f"[Execution] Market order exception: {e}")
            fill = FillResponse(
                status=ResponseStatus.FAILED,
                order_type="market_entry",
                direction=direction.value,
                symbol=self.config.SYMBOL,
                error=str(e),
            )
            return OrderResult.from_fill(fill)

    def place_stop_loss(
        self,
        direction: Direction,
        size_contracts: float,
        stop_price: float,
    ) -> OrderResult:
        """
        Place a stop-loss order (reduce-only).
        Response is available via result.fill (FillResponse).
        """
        is_buy = direction == Direction.SHORT
        action_str = "SL-SELL" if direction == Direction.LONG else "SL-BUY"

        if self.dry_run or not self._exchange:
            fill = FillResponse.dry_run(
                order_type="stop_loss",
                direction=direction.value,
                symbol=self.config.SYMBOL,
                mock_price=stop_price,
                size=size_contracts,
            )
            logger.info(f"[Execution] {fill.summary()}")
            return OrderResult.from_fill(fill)

        try:
            order_type = {
                "trigger": {
                    "triggerPx": stop_price,
                    "isMarket": True,
                    "tpsl": "sl",
                }
            }
            result = self._exchange.order(
                self.config.SYMBOL, is_buy, size_contracts,
                stop_price, order_type, reduce_only=True,
            )
            logger.debug(f"[Execution] Raw SL response: {result}")
            fill = FillResponse.from_hl_response(
                raw=result,
                order_type="stop_loss",
                direction=direction.value,
                symbol=self.config.SYMBOL,
            )
            logger.info(f"[Execution] {fill.summary()}")
            return OrderResult.from_fill(fill)

        except Exception as e:
            logger.error(f"[Execution] Stop-loss placement failed: {e}")
            fill = FillResponse(
                status=ResponseStatus.FAILED,
                order_type="stop_loss",
                direction=direction.value,
                symbol=self.config.SYMBOL,
                error=str(e),
            )
            return OrderResult.from_fill(fill)

    def place_take_profit(
        self,
        direction: Direction,
        size_contracts: float,
        tp_price: float,
    ) -> OrderResult:
        """
        Place a take-profit limit order (reduce-only).
        Response is available via result.fill (FillResponse).
        """
        is_buy = direction == Direction.SHORT
        action_str = "TP-SELL" if direction == Direction.LONG else "TP-BUY"

        if self.dry_run or not self._exchange:
            fill = FillResponse.dry_run(
                order_type="take_profit",
                direction=direction.value,
                symbol=self.config.SYMBOL,
                mock_price=tp_price,
                size=size_contracts,
            )
            logger.info(f"[Execution] {fill.summary()}")
            return OrderResult.from_fill(fill)

        try:
            order_type = {
                "trigger": {
                    "triggerPx": tp_price,
                    "isMarket": False,
                    "tpsl": "tp",
                }
            }
            result = self._exchange.order(
                self.config.SYMBOL, is_buy, size_contracts,
                tp_price, order_type, reduce_only=True,
            )
            logger.debug(f"[Execution] Raw TP response: {result}")
            fill = FillResponse.from_hl_response(
                raw=result,
                order_type="take_profit",
                direction=direction.value,
                symbol=self.config.SYMBOL,
            )
            logger.info(f"[Execution] {fill.summary()}")
            return OrderResult.from_fill(fill)

        except Exception as e:
            logger.error(f"[Execution] Take-profit placement failed: {e}")
            fill = FillResponse(
                status=ResponseStatus.FAILED,
                order_type="take_profit",
                direction=direction.value,
                symbol=self.config.SYMBOL,
                error=str(e),
            )
            return OrderResult.from_fill(fill)

    def cancel_all_orders(self) -> bool:
        """Cancel all open orders for the trading symbol."""
        if self.dry_run or not self._exchange:
            logger.info(f"[Execution] [DRY RUN] Cancel all orders for {self.config.SYMBOL}")
            return True

        try:
            open_orders = self._info.open_orders(self._wallet_address)
            symbol_orders = [o for o in open_orders if o.get("coin") == self.config.SYMBOL]

            if not symbol_orders:
                logger.debug("[Execution] No open orders to cancel.")
                return True

            cancels = [{"coin": o["coin"], "oid": o["oid"]} for o in symbol_orders]
            result = self._exchange.bulk_cancel(cancels)
            logger.info(f"[Execution] Cancelled {len(cancels)} orders: {result}")
            return True
        except Exception as e:
            logger.error(f"[Execution] Cancel all orders failed: {e}")
            return False

    def close_position(self, direction: Direction, size_contracts: float) -> OrderResult:
        """Close an existing position with a market order."""
        close_direction = Direction.SHORT if direction == Direction.LONG else Direction.LONG
        logger.info(f"[Execution] Closing {direction.value} position: {size_contracts:.6f} contracts")
        return self.place_market_order(close_direction, size_contracts)