# mapped_strategy.py

import backtrader as bt
import logging
from typing import Dict

import eventlet
import pandas as pd

from backtest_strategy_instance import BacktestStrategyInstance  # Ensure correct import path

logger = logging.getLogger(__name__)

class MappedStrategy(bt.Strategy):
    """
    A Backtrader Strategy that integrates with a custom StrategyInstance for executing strategy logic.
    """

    params = (
        ('strategy_instance', None),  # Instance of BacktestStrategyInstance
        ('precomputed_indicators', None),  # Dict of precomputed indicators
        ('precomputed_signals', None),  # Dict of precomputed signal states
        ('socketio', None),  # SocketIO instance for emitting progress
        ('socket_conn_id', None),  # Socket connection ID for emitting progress
        ('data_length', None),  # Total number of data points for progress calculation
        ('backtest_name', None)  # Name of the backtest_name
    )

    def __init__(self):
        super().__init__()

        if not self.p.strategy_instance:
            raise ValueError("StrategyInstance must be provided to MappedStrategy.")

        self.strategy_instance: BacktestStrategyInstance = self.p.strategy_instance
        logger.debug(f"StrategyInstance '{self.strategy_instance.strategy_instance_id}' attached to MappedStrategy.")

        # Establish backreference and initialize broker-dependent attributes
        self.strategy_instance.set_backtrader_strategy(self)

        # Now that backtrader_strategy is set, you can safely proceed
        self.precomputed_indicators: Dict[str, pd.DataFrame] = self.p.precomputed_indicators or {}
        self.indicator_pointers: Dict[str, int] = {name: 0 for name in self.precomputed_indicators.keys()}
        self.indicator_names = list(self.precomputed_indicators.keys())

        # Initialize signal data
        self.precomputed_signals: Dict[str, pd.DataFrame] = self.p.precomputed_signals or {}
        self.signal_pointers: Dict[str, int] = {name: 0 for name in self.precomputed_signals.keys()}
        self.signal_names = list(self.precomputed_signals.keys())

        self.current_step = 0

        # Initialize lists to store orders and trades
        self.orders = []
        self.trade_list = []

        # Initialize other needed variables
        self.starting_balance = self.broker.getvalue()
        self.last_progress = 0  # Initialize last_progress
        self.backtest_name = self.p.backtest_name

        self.bar_executed = 0  # Initialize bar_executed
        self.open_trades = {}  # Initialize a dictionary to store open trades

    def notify_order(self, order):
        """
        Handle order notifications from Backtrader.
        """
        logger.debug(f"notify_order called for order {order.ref}, Status: {order.getstatusname()}")

        if order.status in [order.Submitted, order.Accepted]:
            # Order has been submitted/accepted by broker - nothing to do
            return

        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f"BUY EXECUTED, Price: {order.executed.price}, Size: {order.executed.size}")
            elif order.issell():
                self.log(f"SELL EXECUTED, Price: {order.executed.price}, Size: {order.executed.size}")
            self.bar_executed = len(self.datas[0])
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log(f"Order {order.ref} {order.getstatusname()} - {order.info.get('rejectreason', '')}")
            if order.info:
                self.log(f"Order info: {order.info}")
            else:
                self.log(f"No additional info for order {order.ref}.")

        # Remove the order from the list
        if order in self.orders:
            self.orders.remove(order)

    def notify_trade(self, trade):
        logger.debug(
            f"notify_trade called for trade {trade.ref}, PnL: {trade.pnl}, Status: {trade.status_names[trade.status]}")

        if trade.isopen:
            # Trade just opened - use current bar's datetime from data feed
            current_dt = self.data.datetime.datetime(0)
            # Append 'Z' to indicate UTC timezone (Backtrader uses UTC internally)
            open_datetime = (current_dt.isoformat() + 'Z') if current_dt else None
            # Debug logging
            raw_dt = self.data.datetime[0]
            logger.info(f"[DEBUG] Trade open - raw datetime[0]={raw_dt}, converted={current_dt}, iso={open_datetime}")
            self.log(f"TRADE OPENED, Size: {trade.size}, Price: {trade.price}")
            trade_info = {
                'ref': trade.ref,
                'size': trade.size,
                'open_price': trade.price,
                'open_datetime': open_datetime
            }
            # Store the trade_info with trade.ref as key
            self.open_trades[trade.ref] = trade_info
        elif trade.isclosed:
            # Trade just closed - use current bar's datetime from data feed
            current_dt = self.data.datetime.datetime(0)
            # Append 'Z' to indicate UTC timezone (Backtrader uses UTC internally)
            close_datetime = (current_dt.isoformat() + 'Z') if current_dt else None
            self.log(f"TRADE CLOSED, GROSS P/L: {trade.pnl}, NET P/L: {trade.pnlcomm}")
            # Retrieve open trade details
            trade_info = self.open_trades.pop(trade.ref, {})
            # Get the close price from data feed
            close_price = self.data.close[0]
            trade_info.update({
                'close_price': close_price,
                'close_datetime': close_datetime,
                'pnl': trade.pnl,
                'pnlcomm': trade.pnlcomm
            })
            self.trade_list.append(trade_info)

        # Delegate to StrategyInstance if needed
        # self.strategy_instance.notify_trade(trade)

    def log(self, txt, dt=None):
        """Logging function for this strategy"""
        dt = dt or self.datas[0].datetime.datetime(0)
        logger.info(f"{dt.isoformat()} - {txt}")
        # self.strategy_instance.log(txt, dt)

    def next(self):
        # Process any pending order requests if applicable
        # self.process_order_requests()

        self.current_step += 1

        # Debug: Log current price and indicator values every N steps
        if self.current_step <= 10 or self.current_step % 50 == 0:
            current_price = self.data.close[0]
            logger.info(f"[BACKTEST STEP {self.current_step}] Price: {current_price:.2f}")

            # Log all indicator values at this step
            for ind_name, df in self.precomputed_indicators.items():
                idx = self.indicator_pointers.get(ind_name, 0)
                if idx < len(df):
                    row = df.iloc[idx]
                    # Log all columns except 'time'
                    values = {col: row[col] for col in df.columns if col != 'time'}
                    logger.info(f"[BACKTEST STEP {self.current_step}] Indicator '{ind_name}' at idx {idx}: {values}")

        # Execute the strategy logic
        self.execute_strategy()

        # Advance indicator pointers for the next candle
        for name in self.indicator_names:
            if name in self.indicator_pointers:
                self.indicator_pointers[name] += 1

        # Advance signal pointers for the next candle
        for name in self.signal_names:
            if name in self.signal_pointers:
                self.signal_pointers[name] += 1

        # Check if we're at the second-to-last bar
        if self.current_step == (self.p.data_length - 1):
            if self.position:
                self.log(f"Closing remaining position at the second-to-last bar.")
                self.close()

        # Update progress
        if self.p.data_length:
            self.update_progress()

        # Periodically yield to eventlet
        eventlet.sleep(0)

    def execute_strategy(self):
        try:
            execution_result = self.strategy_instance.execute()
            if not execution_result.get('success', False):
                error_msg = execution_result.get('message', 'Unknown error during strategy execution.')
                logger.error(f"Strategy execution failed: {error_msg}")
                self.stop()
        except Exception as e:
            logger.error(f"Error in strategy execution: {e}")

    def update_progress(self):
        progress = (self.current_step / self.p.data_length) * 100
        progress = min(int(progress), 100)

        if progress > self.last_progress:
            if self.p.socketio and self.p.socket_conn_id:
                self.p.socketio.emit(
                    'message',
                    {'reply': 'progress', 'data': { 'test_id': self.backtest_name,
                                                         'progress': progress}},
                    room=self.p.socket_conn_id
                )
                logger.debug(f"Emitted progress: {progress}%")
            self.last_progress = progress

    def stop(self):
        # Close all open positions
        if self.position:
            self.close()
            self.log(f"Closing remaining position at the end of backtest.")

    def t_order(
            self,
            trade_type: str,
            size: float,
            exectype: bt.Order.ExecTypes,
            price: float,
            symbol: str,
            stop_loss_price: float = None,
            take_profit_price: float = None,
            tif: str = 'GTC',
            order_type: str = 'MARKET'
    ):
        """
        Places bracket orders within the strategy context.
        """
        logger.debug(f"t_order called with parameters: {locals()}")

        # Determine the main order method
        if trade_type.lower() == 'buy':
            action = "BUY"
            order_method = self.buy_bracket
        elif trade_type.lower() == 'sell':
            action = "SELL"
            order_method = self.sell_bracket
        else:
            logger.error(f"Invalid trade_type '{trade_type}'. Order not executed.")
            return

        # Prepare bracket order parameters
        bracket_params = {
            'size': size,
            'exectype': exectype,
            'price': price,
            'stopprice': stop_loss_price,
            'limitprice': take_profit_price
        }

        # Remove None values to avoid errors
        bracket_params = {k: v for k, v in bracket_params.items() if v is not None}

        # Place the bracket order
        orders = order_method(**bracket_params)

        # Log the order placement
        message = f"{action} bracket order placed for {size} {symbol} at {order_type} price."
        self.log(message)
        logger.info(message)

        # Store orders
        if not hasattr(self, 'orders'):
            self.orders = []
        self.orders.extend(orders)
