controller.py

"""
controller.py
-------------
Main agent loop for the Momentum Flip trading strategy.

Orchestrates:
1. DataFeed  → live candle data.
2. Strategy  → momentum flip signal evaluation.
3. Risk      → pre-trade risk gate.
4. Execution → order placement.
5. Logger    → structured event logging.

Lifecycle:
- start()  → seeds data, starts WebSocket, enters main loop.
- stop()   → gracefully shuts down all components.
- The loop runs every POLL_INTERVAL_SECONDS seconds.
"""

import time
import signal
import uuid
import logging
from datetime import datetime, timezone
from typing import Optional

from config import Config
from data_feed import DataFeed
from strategy_flip import MomentumFlipStrategy, Direction, SignalResult
from risk import RiskManager
from execution import ExecutionEngine
from logging_agent import AgentLogger

logger = logging.getLogger(__name__)


class ActiveTrade:
    """Tracks a single open trade's lifecycle."""

    def __init__(
        self,
        trade_id: str,
        direction: Direction,
        entry_price: float,
        size_contracts: float,
        sl_price: float,
        tp_price: float,
    ):
        self.trade_id = trade_id
        self.direction = direction
        self.entry_price = entry_price
        self.size_contracts = size_contracts
        self.sl_price = sl_price
        self.tp_price = tp_price
        self.opened_at = datetime.now(timezone.utc)
        self.sl_order_id: Optional[str] = None
        self.tp_order_id: Optional[str] = None


class MomentumFlipController:
    """
    Main controller for the Momentum Flip agent.
    Runs the strategy loop and manages trade lifecycle.
    """

    def __init__(self, config: Config):
        self.config = config
        self.agent_logger = AgentLogger(config)
        self.data_feed = DataFeed(config)
        self.strategy = MomentumFlipStrategy(config)
        self.risk = RiskManager(config)
        self.execution = ExecutionEngine(config, data_feed=self.data_feed)

        self._running = False
        self._active_trades: dict[str, ActiveTrade] = {}
        self._loop_count = 0

        # Register OS signal handlers for graceful shutdown
        signal.signal(signal.SIGINT, self._handle_shutdown)
        signal.signal(signal.SIGTERM, self._handle_shutdown)

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def start(self) -> None:
        """Start the agent: initialize components and enter main loop."""
        self.agent_logger.log_system("AGENT_START", {
            "symbol": self.config.SYMBOL,
            "dry_run": self.config.DRY_RUN,
            "primary_tf": self.config.PRIMARY_TF,
            "filter_tf": self.config.FILTER_TF,
            "entry_tf": self.config.ENTRY_TF,
        })

        logger.info("=" * 60)
        logger.info("  Momentum Flip Agent — Starting")
        logger.info(f"  Symbol:   {self.config.SYMBOL}")
        logger.info(f"  Mode:     {'DRY RUN' if self.config.DRY_RUN else 'LIVE'}")
        logger.info(f"  TFs:      {self.config.ENTRY_TF}m / {self.config.PRIMARY_TF}m / {self.config.FILTER_TF}m")
        logger.info("=" * 60)

        # Set leverage
        self.execution.set_leverage(self.config.LEVERAGE)

        # Start data feed (REST seed + WebSocket)
        self.data_feed.start()

        # Wait for buffers to populate
        logger.info("[Controller] Waiting 3s for candle buffers to populate...")
        time.sleep(3)

        self._running = True
        self._main_loop()

    def stop(self) -> None:
        """Gracefully stop the agent."""
        logger.info("[Controller] Stopping agent...")
        self._running = False
        self.data_feed.stop()
        self.agent_logger.log_system("AGENT_STOP", {
            "active_trades": len(self._active_trades),
            "loop_count": self._loop_count,
        })
        logger.info("[Controller] Agent stopped.")

    def _handle_shutdown(self, signum, frame) -> None:
        """Handle SIGINT/SIGTERM for graceful shutdown."""
        logger.warning(f"[Controller] Received signal {signum} — initiating graceful shutdown.")
        self.stop()

    # ------------------------------------------------------------------
    # Main Loop
    # ------------------------------------------------------------------

    def _main_loop(self) -> None:
        """
        Core agent loop. Runs every POLL_INTERVAL_SECONDS.
        Each iteration:
        1. Fetch current candles.
        2. Evaluate strategy signal.
        3. Check risk.
        4. Execute trade if approved.
        5. Monitor active trades.
        """
        logger.info("[Controller] Main loop started.")

        while self._running:
            loop_start = time.time()
            self._loop_count += 1

            try:
                self._tick()
            except Exception as e:
                self.agent_logger.log_error("Main loop tick", e)
                logger.error(f"[Controller] Unhandled error in tick #{self._loop_count}: {e}")

            # Sleep for remainder of poll interval
            elapsed = time.time() - loop_start
            sleep_time = max(0, self.config.POLL_INTERVAL_SECONDS - elapsed)
            time.sleep(sleep_time)

        logger.info("[Controller] Main loop exited.")

    def _tick(self) -> None:
        """Single iteration of the main loop."""
        # 1. Get candle data for all timeframes
        primary_df = self.data_feed.get_candles(self.config.PRIMARY_TF)
        filter_df  = self.data_feed.get_candles(self.config.FILTER_TF)
        entry_df   = self.data_feed.get_candles(self.config.ENTRY_TF)

        if primary_df.empty or filter_df.empty:
            logger.debug(f"[Controller] Tick #{self._loop_count}: Candle data not ready yet.")
            return

        # 2. Evaluate strategy signal
        signal = self.strategy.evaluate(primary_df, filter_df, entry_df)
        self.agent_logger.log_signal(signal)

        # Heartbeat: show current state every tick so terminal isn't silent
        current_price = self.data_feed.get_latest_price()
        logger.info(
            f"[Tick #{self._loop_count}] "
            f"BTC=${current_price:,.2f} | "
            f"Signal={signal.direction.value} | "
            f"RSI={signal.rsi_value:.1f} | "
            f"MACD_H={signal.macd_histogram:.4f} | "
            f"Confidence={signal.confidence}"
        )

        # 3. Update risk manager with current position count
        open_count = self.execution.get_open_position_count()
        self.risk.set_open_positions(open_count)

        # 4. If signal is actionable, run risk check and execute
        if signal.is_actionable():
            account_equity = self.execution.get_account_equity()
            risk_result = self.risk.check(signal, account_equity)
            self.agent_logger.log_risk_check(risk_result, signal)

            if risk_result.approved:
                self._execute_trade(signal, risk_result)

        # 5. Monitor active trades for SL/TP hits
        self._monitor_active_trades()

        # 6. Log periodic stats every 60 ticks
        if self._loop_count % 60 == 0:
            self.agent_logger.log_daily_stats(self.risk.get_status())

    # ------------------------------------------------------------------
    # Trade Monitoring
    # ------------------------------------------------------------------

    def _monitor_active_trades(self) -> None:
        """
        Check each active trade against the current market price.
        Close the trade if SL or TP has been breached.
        Called every tick so exits are detected promptly.
        """
        if not self._active_trades:
            return

        current_price = self.data_feed.get_latest_price()
        if current_price is None:
            return

        closed_ids = []
        for trade_id, trade in self._active_trades.items():
            reason = None

            if trade.direction == Direction.LONG:
                if current_price <= trade.sl_price:
                    reason = "STOP_LOSS"
                elif current_price >= trade.tp_price:
                    reason = "TAKE_PROFIT"
            else:  # SHORT
                if current_price >= trade.sl_price:
                    reason = "STOP_LOSS"
                elif current_price <= trade.tp_price:
                    reason = "TAKE_PROFIT"

            if reason:
                logger.info(
                    f"[Controller] {reason} hit for {trade_id} "
                    f"@ ${current_price:,.4f} (SL={trade.sl_price} TP={trade.tp_price})"
                )
                closed_ids.append((trade_id, current_price, reason))

        for trade_id, exit_price, reason in closed_ids:
            self.close_trade(trade_id, exit_price, reason)

    # ------------------------------------------------------------------
    # Trade Execution
    # ------------------------------------------------------------------

    def _execute_trade(self, signal: SignalResult, risk_result) -> None:
        """
        Execute a full trade: market entry + SL + TP orders.
        """
        trade_id = f"MF-{uuid.uuid4().hex[:8].upper()}"
        logger.info(f"[Controller] Executing trade {trade_id}: {signal.direction.value}")

        # 1. Place market entry order
        entry_result = self.execution.place_market_order(
            direction=signal.direction,
            size_contracts=risk_result.position_size_contracts,
        )

        self.agent_logger.log_order(
            order_type="MARKET_ENTRY",
            direction=signal.direction.value,
            size=risk_result.position_size_contracts,
            price=entry_result.filled_price,
            success=entry_result.success,
            order_id=entry_result.order_id,
            error=entry_result.error,
        )

        if not entry_result.success:
            logger.error(f"[Controller] Entry order failed for {trade_id}: {entry_result.error}")
            return

        actual_entry = entry_result.filled_price or signal.entry_price

        # 2. Place stop-loss
        sl_result = self.execution.place_stop_loss(
            direction=signal.direction,
            size_contracts=risk_result.position_size_contracts,
            stop_price=risk_result.adjusted_sl,
        )
        self.agent_logger.log_order(
            order_type="STOP_LOSS",
            direction=signal.direction.value,
            size=risk_result.position_size_contracts,
            price=risk_result.adjusted_sl,
            success=sl_result.success,
            order_id=sl_result.order_id,
            error=sl_result.error,
        )

        # 3. Place take-profit
        tp_result = self.execution.place_take_profit(
            direction=signal.direction,
            size_contracts=risk_result.position_size_contracts,
            tp_price=risk_result.adjusted_tp,
        )
        self.agent_logger.log_order(
            order_type="TAKE_PROFIT",
            direction=signal.direction.value,
            size=risk_result.position_size_contracts,
            price=risk_result.adjusted_tp,
            success=tp_result.success,
            order_id=tp_result.order_id,
            error=tp_result.error,
        )

        # 4. Register active trade
        trade = ActiveTrade(
            trade_id=trade_id,
            direction=signal.direction,
            entry_price=actual_entry,
            size_contracts=risk_result.position_size_contracts,
            sl_price=risk_result.adjusted_sl,
            tp_price=risk_result.adjusted_tp,
        )
        trade.sl_order_id = sl_result.order_id
        trade.tp_order_id = tp_result.order_id
        self._active_trades[trade_id] = trade

        self.agent_logger.log_trade_open(
            trade_id=trade_id,
            direction=signal.direction.value,
            symbol=self.config.SYMBOL,
            entry_price=actual_entry,
            size=risk_result.position_size_contracts,
            sl=risk_result.adjusted_sl,
            tp=risk_result.adjusted_tp,
        )

    def close_trade(self, trade_id: str, exit_price: float, reason: str) -> None:
        """
        Close a tracked trade and record PnL.
        Called externally (e.g., from MCP server or monitoring loop).
        """
        trade = self._active_trades.get(trade_id)
        if not trade:
            logger.warning(f"[Controller] Trade {trade_id} not found in active trades.")
            return

        # Calculate PnL
        if trade.direction == Direction.LONG:
            pnl = (exit_price - trade.entry_price) * trade.size_contracts
        else:
            pnl = (trade.entry_price - exit_price) * trade.size_contracts

        self.risk.record_trade_result(pnl)
        self.agent_logger.log_trade_close(
            trade_id=trade_id,
            direction=trade.direction.value,
            symbol=self.config.SYMBOL,
            entry_price=trade.entry_price,
            exit_price=exit_price,
            size=trade.size_contracts,
            pnl=pnl,
            close_reason=reason,
        )

        del self._active_trades[trade_id]

    # ------------------------------------------------------------------
    # Status
    # ------------------------------------------------------------------

    def get_status(self) -> dict:
        """Return current agent status for MCP server queries."""
        return {
            "running": self._running,
            "symbol": self.config.SYMBOL,
            "dry_run": self.config.DRY_RUN,
            "loop_count": self._loop_count,
            "active_trades": len(self._active_trades),
            "active_trade_ids": list(self._active_trades.keys()),
            "risk": self.risk.get_status(),
            "last_signal": (
                {
                    "direction": self.strategy.get_last_signal().direction.value,
                    "confidence": self.strategy.get_last_signal().confidence,
                    "timestamp": str(self.strategy.get_last_signal().timestamp),
                }
                if self.strategy.get_last_signal()
                else None
            ),
        }


# ------------------------------------------------------------------
# Entry Point
# ------------------------------------------------------------------

if __name__ == "__main__":
    import logging
    logging.basicConfig(level=logging.INFO)

    config = Config.from_env()
    controller = MomentumFlipController(config)

    try:
        controller.start()
    except KeyboardInterrupt:
        controller.stop()