engine.py

"""Main trading engine: wires DataFeed, strategy, OMS, executor, and persistence (Chapter 9)."""

from __future__ import annotations

import signal
import threading
import time
from datetime import datetime, timezone
from typing import List, Optional

from loguru import logger

from config import Config
from logging_config import log_trade
from data_feed import CandleBuffer, DataFeed
from executors.paper_executor import PaperExecutor
from indicators import RSIIndicator, TRIXIndicator
from microstructure import TradeEvent
from models.candle import Candle
from models.order_book import OrderBookSnapshot
from oms import OMS
from strategies.base_strategy import (
    BaseStrategy,
    CandleCloseResult,
    ExitSignal,
    IndicatorSnapshot,
    Signal,
    StrategyPosition,
)
from strategies.strategy_02 import Strategy02
from strategies.strategy_03 import Strategy03
from system_state import publish_from_engine, set_engine_running
from trade_db import DEFAULT_SQLITE_PATH, TradeDB


class TradingEngine:
    """Orchestrates tick exits and candle-close entries; blocks in :meth:`start` until :meth:`stop`."""

    def __init__(self, config: Config) -> None:
        self._cfg = config
        self._lock = threading.Lock()
        self._running = False
        self._stop_done = False
        self._last_tick_price: Optional[float] = None
        self._last_upnl_log_m: float = 0.0

        self._db = TradeDB(DEFAULT_SQLITE_PATH, strategy_name=config.STRATEGY_NAME)
        self._oms = OMS(config=config, db=self._db)

        self._strategy = self._build_strategy(config)

        if config.MODE == "MOCK":
            self._executor = PaperExecutor(config=config)
        else:
            from executors.live_executor import LiveExecutor

            self._executor = LiveExecutor(config=config)

        self._rsi = RSIIndicator(length=config.RSI_LEN)
        self._trix = TRIXIndicator(length=config.TRIX_LEN)
        self._buffer = CandleBuffer(maxlen=200)
        self._bar_rsi_prev: Optional[float] = None
        self._bar_trix_prev: Optional[float] = None
        self._dash_rsi: Optional[float] = None
        self._dash_trix: Optional[float] = None
        self._dash_candle: Optional[Candle] = None
        self._dash_micro: Optional[dict] = None

        self._feed = DataFeed(
            config,
            self._on_tick,
            self._on_candle_close,
            on_order_book=self._on_order_book if callable(getattr(self._strategy, "on_book_update", None)) else None,
            on_trade=self._on_trade_event if callable(getattr(self._strategy, "on_trade_event", None)) else None,
        )

        self._log_startup_banner()

    def _build_strategy(self, config: Config) -> BaseStrategy:
        name = (config.STRATEGY_NAME or "strategy_02").strip().lower()
        if name == "strategy_02":
            return Strategy02(config=config)
        if name == "strategy_03":
            return Strategy03(config=config)
        logger.warning("Unknown STRATEGY_NAME='{}' -> falling back to strategy_02", config.STRATEGY_NAME)
        return Strategy02(config=config)

    @staticmethod
    def _candle_to_dict(candle: Candle) -> dict:
        return {
            "open_time": candle.open_time.isoformat(),
            "open": float(candle.open),
            "high": float(candle.high),
            "low": float(candle.low),
            "close": float(candle.close),
        }

    def _publish_dashboard_state(self, tick_ts: Optional[datetime]) -> None:
        """Push latest tick/OHLC/indicators/open legs to the monitoring dashboard."""
        px: Optional[float] = None
        if self._last_tick_price is not None:
            px = float(self._last_tick_price)
        elif self._dash_candle is not None:
            px = float(self._dash_candle.close)
        elif len(self._buffer) > 0:
            px = float(self._buffer.get_candle(0).close)

        ts = tick_ts if tick_ts is not None else datetime.now(timezone.utc)

        pos_dicts: List[dict] = []
        unreal = 0.0
        if px is not None:
            for p in self._oms.get_open_positions():
                unreal += float(p.pnl_usdt)
                pnl_pct = ((px - float(p.entry_price)) / float(p.entry_price) * 100.0) if p.entry_price else 0.0
                pos_dicts.append(
                    {
                        "id": p.id,
                        "symbol": p.symbol,
                        "direction": p.direction,
                        "entry_price": float(p.entry_price),
                        "current_price": px,
                        "unrealized_pnl_usdt": float(p.pnl_usdt),
                        "pnl_pct": pnl_pct,
                        "stop_loss_price": float(p.stop_loss_price),
                        "entry_reason": p.entry_reason or "",
                    }
                )
        else:
            for p in self._oms.get_open_positions():
                unreal += float(p.pnl_usdt)
                mark = float(p.mark_price or p.entry_price)
                pnl_pct = ((mark - float(p.entry_price)) / float(p.entry_price) * 100.0) if p.entry_price else 0.0
                pos_dicts.append(
                    {
                        "id": p.id,
                        "symbol": p.symbol,
                        "direction": p.direction,
                        "entry_price": float(p.entry_price),
                        "current_price": mark,
                        "unrealized_pnl_usdt": float(p.pnl_usdt),
                        "pnl_pct": pnl_pct,
                        "stop_loss_price": float(p.stop_loss_price),
                        "entry_reason": p.entry_reason or "",
                    }
                )

        paper_eq = None
        if isinstance(self._executor, PaperExecutor):
            paper_eq = float(self._executor.get_simulated_balance()["equity"])

        candle_d = self._candle_to_dict(self._dash_candle) if self._dash_candle else None

        publish_from_engine(
            price=px,
            tick_utc=ts,
            rsi=self._dash_rsi,
            trix=self._dash_trix,
            candle=candle_d,
            positions=pos_dicts,
            daily_trade_count=self._oms.get_daily_trade_count(),
            unrealized_total=unreal,
            paper_equity=paper_eq,
            micro=self._dash_micro,
        )

    def _log_startup_banner(self) -> None:
        logger.info(
            "======== TradingEngine ========\n"
            "  symbol={} | timeframe={} | mode={} | session_end_utc={}\n"
            "  risk: max_trades_in_row={} | max_open_positions={} | max_trades_per_day={} | "
            "trade_qty={}\n"
            "  exits: long_sl={} long_trail={} long_th={} | short_sl={} short_trail={} short_th={}\n"
            "  indicators: rsi_len={} rsi_ob/os={}/{} | trix_len={} | tick_size={}\n"
            "================================",
            self._cfg.SYMBOL,
            self._cfg.TIMEFRAME,
            self._cfg.MODE,
            self._cfg.SESSION_END_UTC,
            self._cfg.MAX_TRADES_IN_ROW,
            self._cfg.MAX_OPEN_POSITIONS,
            self._cfg.MAX_TRADES_PER_DAY,
            self._cfg.TRADE_QTY,
            self._cfg.LONG_SL,
            self._cfg.LONG_TRAIL,
            self._cfg.LONG_TH,
            self._cfg.SHORT_SL,
            self._cfg.SHORT_TRAIL,
            self._cfg.SHORT_TH,
            self._cfg.RSI_LEN,
            self._cfg.RSI_OVERBOUGHT,
            self._cfg.RSI_OVERSOLD,
            self._cfg.TRIX_LEN,
            self._cfg.TICK_SIZE,
        )

    def _positions_for_strategy(self) -> List[StrategyPosition]:
        return [
            StrategyPosition(id=p.id, direction=p.direction, entry_price=float(p.entry_price))
            for p in self._oms.get_open_positions()
        ]

    def _flatten_price(self) -> Optional[float]:
        if self._last_tick_price is not None:
            return float(self._last_tick_price)
        if len(self._buffer) > 0:
            return float(self._buffer.get_candle(0).close)
        return None

    def _maybe_log_unrealized(self, open_count: int) -> None:
        now = time.monotonic()
        if now - self._last_upnl_log_m < 5.0:
            return
        self._last_upnl_log_m = now
        if isinstance(self._executor, PaperExecutor):
            bal = self._executor.get_simulated_balance()
            logger.info(
                "uPnL | unrealized={:.6f} | equity={:.6f} | open_legs={}",
                bal["unrealized_pnl"],
                bal["equity"],
                open_count,
            )
        else:
            positions = self._oms.get_open_positions()
            total = sum(p.pnl_usdt for p in positions)
            logger.info("uPnL | unrealized_sum={:.6f} | open_legs={}", total, len(positions))

    def _process_exit(self, ex: ExitSignal) -> None:
        ids = {p.id for p in self._oms.get_open_positions()}
        if ex.position_id not in ids:
            logger.warning("ExitSignal unknown or already flat | id={}", ex.position_id)
            return
        pos = next(p for p in self._oms.get_open_positions() if p.id == ex.position_id)
        res = self._executor.execute_exit(pos, ex)
        if not res.filled:
            logger.error("Exit not filled | id={} | {}", ex.position_id, res.message)
            return
        filled = ExitSignal(
            position_id=ex.position_id,
            reason=ex.reason,
            price=float(res.fill_price),
        )
        self._oms.close_position(ex.position_id, filled)
        sync_exit = getattr(self._strategy, "sync_position_after_exit", None)
        if callable(sync_exit):
            sync_exit(ex.position_id)

    def _process_entry(self, sig: Signal, rsi_v: Optional[float], trix_v: Optional[float]) -> None:
        allowed, reason = self._oms.validate_entry(sig)
        if not allowed:
            self._db.save_signal(
                sig,
                False,
                reason,
                rsi_value=rsi_v,
                trix_value=trix_v,
            )
            logger.info(
                "SIGNAL | {} {} @ {:.8f} | BLOCKED | {}",
                sig.direction,
                sig.reason,
                sig.price,
                reason,
            )
            return

        # Get ATR from strategy for position sizing
        atr = None
        get_ms = getattr(self._strategy, "get_microstructure_snapshot", None)
        if callable(get_ms):
            snap = get_ms()
            atr = snap.atr if snap else None

        pos = self._oms.open_position(sig, atr=atr)
        res = self._executor.execute_entry(sig, pos)
        if res.filled:
            fp = float(res.fill_price)
            pos.entry_price = fp
            pos.mark_price = fp
            sync_entry = getattr(self._strategy, "sync_position_after_entry", None)
            if callable(sync_entry):
                sync_entry(pos.id, sig.direction, fp)
            self._db.save_signal(sig, True, "", rsi_value=rsi_v, trix_value=trix_v)
            log_trade(
                "SIGNAL | {} {} @ {:.8f} | ALLOWED | fill={:.8f} | qty={:.6f}",
                sig.direction,
                sig.reason,
                sig.price,
                fp,
                pos.qty,
            )
        else:
            self._oms.abandon_open_position(pos.id)
            msg = res.message or "entry not filled"
            self._db.save_signal(
                sig,
                False,
                msg,
                rsi_value=rsi_v,
                trix_value=trix_v,
            )
            logger.info(
                "SIGNAL | {} {} @ {:.8f} | BLOCKED | {}",
                sig.direction,
                sig.reason,
                sig.price,
                msg,
            )

    def _handle_candle_result(
        self,
        result: CandleCloseResult,
        rsi_v: Optional[float],
        trix_v: Optional[float],
    ) -> None:
        if result.exit_signal is not None:
            self._process_exit(result.exit_signal)
        if result.entry_signal is not None:
            self._process_entry(result.entry_signal, rsi_v, trix_v)

    def _on_order_book(self, snapshot: OrderBookSnapshot) -> None:
        fn = getattr(self._strategy, "on_book_update", None)
        if callable(fn):
            fn(snapshot)

    def _on_trade_event(self, event: TradeEvent) -> None:
        fn = getattr(self._strategy, "on_trade_event", None)
        if callable(fn):
            fn(event)

    def _on_tick(self, price: float, _timestamp: datetime) -> None:
        with self._lock:
            if isinstance(self._executor, PaperExecutor):
                self._executor.update_tick(price)
            self._last_tick_price = float(price)
            self._oms.apply_tick_to_positions(price)
            strat_positions = self._positions_for_strategy()

            # Strategy entry-on-tick hook (implemented by strategy_02/03).
            eval_entry = getattr(self._strategy, "evaluate_entry", None)
            get_ms = getattr(self._strategy, "get_microstructure_snapshot", None)
            if callable(eval_entry) and not strat_positions:
                sig = eval_entry(price)
                ms = get_ms() if callable(get_ms) else None
                if ms is not None:
                    # Publish microstructure metrics for debugging / dashboard.
                    self._dash_micro = {
                        "obi": float(ms.obi),
                        "delta": float(ms.delta),
                        "delta_cumulative": float(ms.delta_cumulative),
                        "absorption_long": bool(ms.absorption_long),
                        "absorption_short": bool(ms.absorption_short),
                        "spread_pct": float(ms.spread_pct),
                        "mid_price": float(ms.mid_price),
                        "atr": (float(ms.atr) if ms.atr is not None else None),
                        "bid_volume": float(ms.bid_volume),
                        "ask_volume": float(ms.ask_volume),
                    }
                if sig is not None:
                    # RSI/TRIX are candle-based and may still be warming up on new starts.
                    # Keep signal metadata consistent: save the actual indicator values.
                    self._process_entry(sig, self._dash_rsi, self._dash_trix)

            ex = self._strategy.on_tick(price, strat_positions)
            if ex is not None:
                self._process_exit(ex)
            self._maybe_log_unrealized(len(self._oms.get_open_positions()))
            self._publish_dashboard_state(_timestamp)

    def _on_candle_close(self, candle: Candle) -> None:
        with self._lock:
            self._buffer.append_closed(candle)
            rsi_new = self._rsi.update(float(candle.close))
            trix_new = self._trix.update(float(candle.close))
            open_n = len(self._oms.get_open_positions())

            ts = candle.open_time.isoformat()
            o, h, l, c = candle.open, candle.high, candle.low, candle.close
            if rsi_new is None or trix_new is None:
                self._dash_candle = candle
                self._dash_rsi = None
                self._dash_trix = None
                logger.info(
                    "BAR | t={} | O={:.8f} H={:.8f} L={:.8f} C={:.8f} | RSI={} TRIX={} | positions={}",
                    ts,
                    o,
                    h,
                    l,
                    c,
                    "—" if rsi_new is None else f"{rsi_new:.4f}",
                    "—" if trix_new is None else f"{trix_new:.6f}",
                    open_n,
                )
                self._publish_dashboard_state(candle.open_time)
                return

            self._dash_candle = candle
            self._dash_rsi = rsi_new
            self._dash_trix = trix_new

            snap = IndicatorSnapshot(
                self._bar_rsi_prev,
                rsi_new,
                self._bar_trix_prev,
                trix_new,
            )
            self._bar_rsi_prev = rsi_new
            self._bar_trix_prev = trix_new

            result = self._strategy.on_candle_close(candle, snap)

            logger.info(
                "BAR | t={} | O={:.8f} H={:.8f} L={:.8f} C={:.8f} | RSI={:.4f} TRIX={:.6f} | positions={}",
                ts,
                o,
                h,
                l,
                c,
                rsi_new,
                trix_new,
                open_n,
            )

            self._handle_candle_result(result, rsi_new, trix_new)
            self._publish_dashboard_state(candle.open_time)

    def start(self) -> None:
        if self._running:
            logger.warning("TradingEngine.start() ignored (already running)")
            return

        self._running = True
        self._stop_done = False
        self._last_upnl_log_m = time.monotonic()

        def _handle_sigint(_signum: int, _frame: object) -> None:
            logger.info("SIGINT received — shutting down")
            self.stop()

        signal.signal(signal.SIGINT, _handle_sigint)

        self._feed.start()
        set_engine_running(True)
        try:
            while not self._stop_done:
                time.sleep(0.2)
        except KeyboardInterrupt:
            logger.info("KeyboardInterrupt — shutting down")
            self.stop()
        finally:
            self.stop()

    def stop(self) -> None:
        set_engine_running(False)
        with self._lock:
            if self._stop_done:
                return
            self._stop_done = True
            self._running = False

        logger.info("TradingEngine stopping — closing feed and flattening")

        self._feed.stop()

        px = self._flatten_price()
        with self._lock:
            opens = list(self._oms.get_open_positions())
        if not opens:
            logger.info("TradingEngine stopped | no open positions")
            return

        if px is None:
            logger.warning(
                "Cannot flatten: no tick/candle price | {} open leg(s) left in OMS",
                len(opens),
            )
            return

        logger.info("Flattening {} position(s) @ {:.8f} (SESSION_END)", len(opens), px)
        for pos in opens:
            ex = ExitSignal(position_id=pos.id, reason="SESSION_END", price=px)
            try:
                self._process_exit(ex)
            except Exception as exc:  # noqa: BLE001
                logger.exception("Flatten failed | id={} | {}", pos.id, exc)

        logger.info("TradingEngine stopped cleanly")