#!/usr/bin/env python3
"""
Live Trading Bot for Regime Reversion Strategy on OKX.

This script runs the regime-based mean reversion strategy
on ETH perpetual futures using OKX exchange.

Usage:
    # Run with demo account (default)
    uv run python -m live_trading.main
    
    # Run with specific settings
    uv run python -m live_trading.main --max-position 500 --leverage 2
    
    # Run without UI (headless mode)
    uv run python -m live_trading.main --no-ui
"""
import argparse
import logging
import queue
import signal
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))

from live_trading.config import get_config, OKXConfig, TradingConfig, PathConfig
from live_trading.okx_client import OKXClient
from live_trading.data_feed import DataFeed
from live_trading.position_manager import PositionManager
from live_trading.live_regime_strategy import LiveRegimeStrategy
from live_trading.db.database import init_db, TradingDatabase
from live_trading.db.migrations import run_migrations
from live_trading.ui.state import SharedState, PositionState
from live_trading.ui.dashboard import TradingDashboard, setup_ui_logging


def setup_logging(
    log_dir: Path,
    log_queue: Optional[queue.Queue] = None,
) -> logging.Logger:
    """
    Configure logging for the trading bot.
    
    Args:
        log_dir: Directory for log files
        log_queue: Optional queue for UI log handler
        
    Returns:
        Logger instance
    """
    log_file = log_dir / "live_trading.log"
    
    handlers = [
        logging.FileHandler(log_file),
    ]
    
    # Only add StreamHandler if no UI (log_queue is None)
    if log_queue is None:
        handlers.append(logging.StreamHandler(sys.stdout))
    
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s [%(levelname)s] %(name)s: %(message)s',
        handlers=handlers,
        force=True,
    )
    
    # Add UI log handler if queue provided
    if log_queue is not None:
        setup_ui_logging(log_queue)
    
    return logging.getLogger(__name__)


class LiveTradingBot:
    """
    Main trading bot orchestrator.
    
    Coordinates data fetching, signal generation, and order execution
    in a continuous loop.
    """
    
    def __init__(
        self,
        okx_config: OKXConfig,
        trading_config: TradingConfig,
        path_config: PathConfig,
        database: Optional[TradingDatabase] = None,
        shared_state: Optional[SharedState] = None,
    ):
        self.okx_config = okx_config
        self.trading_config = trading_config
        self.path_config = path_config
        self.db = database
        self.state = shared_state
        
        self.logger = logging.getLogger(__name__)
        self.running = True
        
        # Initialize components
        self.logger.info("Initializing trading bot components...")
        
        self.okx_client = OKXClient(okx_config, trading_config)
        self.data_feed = DataFeed(self.okx_client, trading_config, path_config)
        self.position_manager = PositionManager(
            self.okx_client, trading_config, path_config, database
        )
        self.strategy = LiveRegimeStrategy(trading_config, path_config)
        
        # Register signal handlers for graceful shutdown
        signal.signal(signal.SIGINT, self._handle_shutdown)
        signal.signal(signal.SIGTERM, self._handle_shutdown)
        
        # Initialize shared state if provided
        if self.state:
            mode = "DEMO" if okx_config.demo_mode else "LIVE"
            self.state.set_mode(mode)
            self.state.set_symbols(
                trading_config.eth_symbol,
                trading_config.btc_symbol,
            )
            self.state.update_account(0.0, 0.0, trading_config.leverage)
        
        self._print_startup_banner()
    
    def _print_startup_banner(self) -> None:
        """Print startup information."""
        mode = "DEMO/SANDBOX" if self.okx_config.demo_mode else "LIVE"
        
        print("=" * 60)
        print(f"  Regime Reversion Strategy - Live Trading Bot")
        print("=" * 60)
        print(f"  Mode:           {mode}")
        print(f"  Trading Pair:   {self.trading_config.eth_symbol}")
        print(f"  Context Pair:   {self.trading_config.btc_symbol}")
        print(f"  Timeframe:      {self.trading_config.timeframe}")
        print(f"  Max Position:   ${self.trading_config.max_position_usdt if self.trading_config.max_position_usdt > 0 else 'All available'}")
        print(f"  Leverage:       {self.trading_config.leverage}x")
        print(f"  Stop Loss:      {self.trading_config.stop_loss_pct * 100:.1f}%")
        print(f"  Take Profit:    {self.trading_config.take_profit_pct * 100:.1f}%")
        print(f"  Cycle Interval: {self.trading_config.sleep_seconds // 60} minutes")
        print("=" * 60)
        
        if not self.okx_config.demo_mode:
            print("\n  *** WARNING: LIVE TRADING MODE - REAL FUNDS AT RISK ***\n")
    
    def _handle_shutdown(self, signum, frame) -> None:
        """Handle shutdown signals gracefully."""
        self.logger.info("Shutdown signal received, stopping...")
        self.running = False
        if self.state:
            self.state.stop()
    
    def run_trading_cycle(self) -> None:
        """
        Execute one trading cycle.
        
        1. Fetch latest market data
        2. Update open positions
        3. Generate trading signal
        4. Execute trades if signal triggers
        5. Update shared state for UI
        """
        # Reload model if it has changed (e.g. daily training)
        try:
            self.strategy.reload_model_if_changed()
        except Exception as e:
            self.logger.warning(f"Failed to reload model: {e}")

        cycle_start = datetime.now(timezone.utc)
        self.logger.info(f"--- Trading Cycle Start: {cycle_start.isoformat()} ---")
        
        if self.state:
            self.state.set_last_cycle_time(cycle_start.isoformat())
        
        try:
            # 1. Fetch market data
            features = self.data_feed.get_latest_data()
            if features is None or features.empty:
                self.logger.warning("No market data available, skipping cycle")
                return
            
            # Get current prices
            eth_price = features['eth_close'].iloc[-1]
            btc_price = features['btc_close'].iloc[-1]
            
            current_prices = {
                self.trading_config.eth_symbol: eth_price,
                self.trading_config.btc_symbol: btc_price,
            }
            
            # 2. Update existing positions (check SL/TP)
            closed_trades = self.position_manager.update_positions(current_prices)
            if closed_trades:
                for trade in closed_trades:
                    self.logger.info(
                        f"Trade closed: {trade['trade_id']} "
                        f"PnL=${trade['pnl_usd']:.2f} ({trade['reason']})"
                    )
            
            # 3. Sync with exchange positions
            self.position_manager.sync_with_exchange()
            
            # 4. Get current funding rates
            funding = self.data_feed.get_current_funding_rates()
            
            # 5. Generate trading signal
            sig = self.strategy.generate_signal(features, funding)
            
            # 6. Update shared state with strategy info
            self._update_strategy_state(sig, funding)
            
            # 7. Execute trades based on signal
            if sig['action'] == 'entry':
                self._execute_entry(sig, eth_price)
            elif sig['action'] == 'check_exit':
                self._execute_exit(sig)
            
            # 8. Update shared state with position and account
            self._update_position_state(eth_price)
            self._update_account_state()
            
            # 9. Log portfolio summary
            summary = self.position_manager.get_portfolio_summary()
            self.logger.info(
                f"Portfolio: {summary['open_positions']} positions, "
                f"exposure=${summary['total_exposure_usdt']:.2f}, "
                f"unrealized_pnl=${summary['total_unrealized_pnl']:.2f}"
            )
            
        except Exception as e:
            self.logger.error(f"Trading cycle error: {e}", exc_info=True)
            # Save positions on error
            self.position_manager.save_positions()
        
        cycle_duration = (datetime.now(timezone.utc) - cycle_start).total_seconds()
        self.logger.info(f"--- Cycle completed in {cycle_duration:.1f}s ---")
    
    def _update_strategy_state(self, sig: dict, funding: dict) -> None:
        """Update shared state with strategy information."""
        if not self.state:
            return
        
        self.state.update_strategy(
            z_score=sig.get('z_score', 0.0),
            probability=sig.get('probability', 0.0),
            funding_rate=funding.get('btc_funding', 0.0),
            action=sig.get('action', 'hold'),
            reason=sig.get('reason', ''),
        )
    
    def _update_position_state(self, current_price: float) -> None:
        """Update shared state with current position."""
        if not self.state:
            return
        
        symbol = self.trading_config.eth_symbol
        position = self.position_manager.get_position_for_symbol(symbol)
        
        if position is None:
            self.state.clear_position()
            return
        
        pos_state = PositionState(
            trade_id=position.trade_id,
            symbol=position.symbol,
            side=position.side,
            entry_price=position.entry_price,
            current_price=position.current_price,
            size=position.size,
            size_usdt=position.size_usdt,
            unrealized_pnl=position.unrealized_pnl,
            unrealized_pnl_pct=position.unrealized_pnl_pct,
            stop_loss_price=position.stop_loss_price,
            take_profit_price=position.take_profit_price,
        )
        self.state.set_position(pos_state)
    
    def _update_account_state(self) -> None:
        """Update shared state with account information."""
        if not self.state:
            return
        
        try:
            balance = self.okx_client.get_balance()
            self.state.update_account(
                balance=balance.get('total', 0.0),
                available=balance.get('free', 0.0),
                leverage=self.trading_config.leverage,
            )
        except Exception as e:
            self.logger.warning(f"Failed to update account state: {e}")
    
    def _execute_entry(self, signal: dict, current_price: float) -> None:
        """Execute entry trade."""
        symbol = self.trading_config.eth_symbol
        side = signal['side']
        
        # Check if we can open a position
        if not self.position_manager.can_open_position():
            self.logger.info("Cannot open position: max positions reached")
            return
        
        # Get account balance
        balance = self.okx_client.get_balance()
        available_usdt = balance['free']
        self.logger.info(f"Account balance: ${available_usdt:.2f} USDT available")

        # Calculate position size
        size_usdt = self.strategy.calculate_position_size(signal, available_usdt)
        if size_usdt <= 0:
            self.logger.info(
                f"Position size too small (${size_usdt:.2f}), skipping entry. "
                f"Min required: ${self.strategy.config.min_position_usdt:.2f}"
            )
            return
        
        size_eth = size_usdt / current_price
        
        # Calculate SL/TP for logging
        stop_loss, take_profit = self.strategy.calculate_sl_tp(current_price, side)
        
        sl_str = f"{stop_loss:.2f}" if stop_loss else "N/A"
        tp_str = f"{take_profit:.2f}" if take_profit else "N/A"
        self.logger.info(
            f"Executing {side.upper()} entry: {size_eth:.4f} ETH @ {current_price:.2f} "
            f"(${size_usdt:.2f}), SL={sl_str}, TP={tp_str}"
        )
        
        try:
            # Place market order (guaranteed to have fill price or raises)
            order_side = "buy" if side == "long" else "sell"
            order = self.okx_client.place_market_order(symbol, order_side, size_eth)
            
            # Get filled price and amount (guaranteed by OKX client)
            filled_price = order['average']
            filled_amount = order.get('filled') or size_eth
            
            # Calculate SL/TP with filled price
            stop_loss, take_profit = self.strategy.calculate_sl_tp(filled_price, side)
            
            if stop_loss is None or take_profit is None:
                raise RuntimeError(
                    f"Failed to calculate SL/TP: filled_price={filled_price}, side={side}"
                )
            
            # Get order ID from response
            order_id = order.get('id', '')
            
            # Record position locally
            position = self.position_manager.open_position(
                symbol=symbol,
                side=side,
                entry_price=filled_price,
                size=filled_amount,
                stop_loss_price=stop_loss,
                take_profit_price=take_profit,
                order_id=order_id,
            )
            
            if position:
                self.logger.info(
                    f"Position opened: {position.trade_id}, "
                    f"{filled_amount:.4f} ETH @ {filled_price:.2f}"
                )
                
                # Try to set SL/TP on exchange
                try:
                    self.okx_client.set_stop_loss_take_profit(
                        symbol, side, filled_amount, stop_loss, take_profit
                    )
                except Exception as e:
                    self.logger.warning(f"Could not set SL/TP on exchange: {e}")
            
        except Exception as e:
            self.logger.error(f"Order execution failed: {e}", exc_info=True)
    
    def _execute_exit(self, signal: dict) -> None:
        """Execute exit based on mean reversion signal."""
        symbol = self.trading_config.eth_symbol
        
        # Get position for ETH
        position = self.position_manager.get_position_for_symbol(symbol)
        if not position:
            return
        
        current_price = signal.get('eth_price', position.current_price)
        
        self.logger.info(
            f"Mean reversion exit signal: closing {position.trade_id} "
            f"@ {current_price:.2f}"
        )
        
        try:
            # Close position on exchange
            exit_order = self.okx_client.close_position(symbol)
            exit_order_id = exit_order.get('id', '') if exit_order else ''
            
            # Record closure locally
            self.position_manager.close_position(
                position.trade_id,
                current_price,
                reason="mean_reversion_complete",
                exit_order_id=exit_order_id,
            )
            
        except Exception as e:
            self.logger.error(f"Exit execution failed: {e}", exc_info=True)
    
    def _is_running(self) -> bool:
        """Check if bot should continue running."""
        if not self.running:
            return False
        if self.state and not self.state.is_running():
            return False
        return True
    
    def run(self) -> None:
        """Main trading loop."""
        self.logger.info("Starting trading loop...")
        
        while self._is_running():
            try:
                self.run_trading_cycle()
                
                if self._is_running():
                    sleep_seconds = self.trading_config.sleep_seconds
                    minutes = sleep_seconds // 60
                    self.logger.info(f"Sleeping for {minutes} minutes...")
                    
                    # Sleep in smaller chunks to allow faster shutdown
                    for _ in range(sleep_seconds):
                        if not self._is_running():
                            break
                        time.sleep(1)
                        
            except KeyboardInterrupt:
                self.logger.info("Keyboard interrupt received")
                break
            except Exception as e:
                self.logger.error(f"Unexpected error in main loop: {e}", exc_info=True)
                time.sleep(60)  # Wait before retry
        
        # Cleanup
        self.logger.info("Shutting down...")
        self.position_manager.save_positions()
        if self.db:
            self.db.close()
        self.logger.info("Shutdown complete")


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description="Live Trading Bot for Regime Reversion Strategy"
    )
    parser.add_argument(
        "--max-position",
        type=float,
        default=None,
        help="Maximum position size in USDT"
    )
    parser.add_argument(
        "--leverage",
        type=int,
        default=None,
        help="Trading leverage (1-125)"
    )
    parser.add_argument(
        "--interval",
        type=int,
        default=None,
        help="Trading cycle interval in seconds"
    )
    parser.add_argument(
        "--live",
        action="store_true",
        help="Use live trading mode (requires OKX_DEMO_MODE=false)"
    )
    parser.add_argument(
        "--no-ui",
        action="store_true",
        help="Run in headless mode without terminal UI"
    )
    return parser.parse_args()


def main():
    """Main entry point."""
    args = parse_args()
    
    # Load configuration
    okx_config, trading_config, path_config = get_config()
    
    # Apply command line overrides
    if args.max_position is not None:
        trading_config.max_position_usdt = args.max_position
    if args.leverage is not None:
        trading_config.leverage = args.leverage
    if args.interval is not None:
        trading_config.sleep_seconds = args.interval
    if args.live:
        okx_config.demo_mode = False
    
    # Determine if UI should be enabled
    use_ui = not args.no_ui and sys.stdin.isatty()
    
    # Initialize database
    db_path = path_config.base_dir / "live_trading" / "trading.db"
    db = init_db(db_path)
    
    # Run migrations (imports CSV if exists)
    run_migrations(db, path_config.trade_log_file)
    
    # Initialize UI components if enabled
    log_queue: Optional[queue.Queue] = None
    shared_state: Optional[SharedState] = None
    dashboard: Optional[TradingDashboard] = None
    
    if use_ui:
        log_queue = queue.Queue(maxsize=1000)
        shared_state = SharedState()
    
    # Setup logging
    logger = setup_logging(path_config.logs_dir, log_queue)
    
    try:
        # Create bot
        bot = LiveTradingBot(
            okx_config,
            trading_config,
            path_config,
            database=db,
            shared_state=shared_state,
        )
        
        # Start dashboard if UI enabled
        if use_ui and shared_state and log_queue:
            dashboard = TradingDashboard(
                state=shared_state,
                db=db,
                log_queue=log_queue,
                on_quit=lambda: setattr(bot, 'running', False),
            )
            dashboard.start()
            logger.info("Dashboard started")
        
        # Run bot
        bot.run()
        
    except ValueError as e:
        logger.error(f"Configuration error: {e}")
        sys.exit(1)
    except Exception as e:
        logger.error(f"Fatal error: {e}", exc_info=True)
        sys.exit(1)
    finally:
        # Cleanup
        if dashboard:
            dashboard.stop()
        if db:
            db.close()


if __name__ == "__main__":
    main()