# backtest_strategy_instance.py

import logging
from typing import Any, Optional

import pandas as pd
import datetime as dt
import backtrader as bt
from StrategyInstance import StrategyInstance

logger = logging.getLogger(__name__)


class BacktestStrategyInstance(StrategyInstance):
    """
    Extends StrategyInstance with custom methods for backtesting.
    """

    def __init__(self, strategy_instance_id: str, strategy_id: str, strategy_name: str,
                 user_id: int, generated_code: str, data_cache: Any, indicators: Any | None,
                 trades: Any | None, backtrader_strategy: Optional[bt.Strategy] = None):
        # Set 'self.broker' and 'self.backtrader_strategy' to None before calling super().__init__()
        self.broker = None
        self.backtrader_strategy = None

        super().__init__(strategy_instance_id, strategy_id, strategy_name, user_id,
                         generated_code, data_cache, indicators, trades)

        # Set the backtrader_strategy instance after super().__init__()
        self.backtrader_strategy = backtrader_strategy
        self.broker = self.backtrader_strategy.broker if self.backtrader_strategy else None

        # Initialize balances; they will be set after backtrader_strategy is available
        self.starting_balance = 0.0
        self.current_balance = 0.0
        self.available_balance = 0.0
        self.available_strategy_balance = 0.0

        # Update exec_context with balance attributes
        self.exec_context['starting_balance'] = self.starting_balance
        self.exec_context['current_balance'] = self.current_balance
        self.exec_context['available_balance'] = self.available_balance
        self.exec_context['available_strategy_balance'] = self.available_strategy_balance

        # Initialize last_valid_values for indicators
        self.last_valid_values={}

    def set_backtrader_strategy(self, backtrader_strategy: bt.Strategy):
        """
        Sets the backtrader_strategy and initializes broker-dependent attributes.
        """
        self.backtrader_strategy = backtrader_strategy
        self.broker = self.backtrader_strategy.broker

        # Now initialize balances from Backtrader's broker
        self.starting_balance = self.fetch_user_balance()
        self.current_balance = self.starting_balance
        self.available_balance = self.calculate_available_balance()
        self.available_strategy_balance = self.starting_balance

        # Update exec_context with updated balance attributes
        self.exec_context['starting_balance'] = self.starting_balance
        self.exec_context['current_balance'] = self.current_balance
        self.exec_context['available_balance'] = self.available_balance
        self.exec_context['available_strategy_balance'] = self.available_strategy_balance

    # 1. Override trade_order
    def trade_order(
            self,
            trade_type: str,
            size: float,
            order_type: str,
            source: dict = None,
            tif: str = 'GTC',
            stop_loss: dict = None,
            trailing_stop: dict = None,
            take_profit: dict = None,
            limit: dict = None,
            trailing_limit: dict = None,
            target_market: dict = None,
            name_order: dict = None
    ):
        """
        Custom trade_order method for backtesting.
        Executes trades within the Backtrader environment and updates balances.
        """
        if self.backtrader_strategy is None:
            logger.error("Backtrader strategy is not set in StrategyInstance.")
            return

        # Validate and extract symbol
        symbol = source.get('symbol') or source.get('market') if source else 'Unknown'
        if not symbol:
            logger.error("Symbol not provided in source. Order not executed.")
            return

        # Get current price from Backtrader's data feed
        price = self.get_current_price()

        # Get stop_loss and take_profit prices
        stop_loss_price = stop_loss.get('value') if stop_loss else None
        take_profit_price = take_profit.get('value') if take_profit else None

        # Determine execution type based on order_type
        order_type = order_type.upper()
        if order_type == 'MARKET':
            exectype = bt.Order.Market
            order_price = None  # Do not set price for market orders
        elif order_type == 'LIMIT':
            exectype = bt.Order.Limit
            order_price = price  # Use current price as the limit price
        else:
            logger.error(f"Invalid order_type '{order_type}'. Order not executed.")
            return

        # Place the main order and associated stop loss and take profit orders
        if trade_type.lower() == 'buy':
            action = "BUY"
            # Place main buy order
            main_order = self.backtrader_strategy.buy(
                size=size,
                price=order_price,
                exectype=exectype,
                transmit=False
            )
        elif trade_type.lower() == 'sell':
            action = "SELL"
            # Place main sell order
            main_order = self.backtrader_strategy.sell(
                size=size,
                price=order_price,
                exectype=exectype,
                transmit=False
            )
        else:
            logger.error(f"Invalid trade_type '{trade_type}'. Order not executed.")
            return

        # Place stop loss order if specified
        if stop_loss_price is not None:
            if trade_type.lower() == 'buy':
                # Stop loss is a sell order for a buy position
                stop_order = self.backtrader_strategy.sell(
                    size=size,
                    price=stop_loss_price,
                    exectype=bt.Order.Stop,
                    parent=main_order,
                    transmit=False
                )
            else:
                # Stop loss is a buy order for a sell position
                stop_order = self.backtrader_strategy.buy(
                    size=size,
                    price=stop_loss_price,
                    exectype=bt.Order.Stop,
                    parent=main_order,
                    transmit=False
                )
        else:
            stop_order = None

        # Place take profit order if specified
        if take_profit_price is not None:
            if trade_type.lower() == 'buy':
                # Take profit is a sell order for a buy position
                limit_order = self.backtrader_strategy.sell(
                    size=size,
                    price=take_profit_price,
                    exectype=bt.Order.Limit,
                    parent=main_order,
                    transmit=False
                )
            else:
                # Take profit is a buy order for a sell position
                limit_order = self.backtrader_strategy.buy(
                    size=size,
                    price=take_profit_price,
                    exectype=bt.Order.Limit,
                    parent=main_order,
                    transmit=False
                )
        else:
            limit_order = None

        # Transmit the entire order chain
        if stop_order and limit_order:
            # Both stop loss and take profit
            limit_order.transmit = True
        elif stop_order or limit_order:
            # Only one of stop loss or take profit
            (stop_order or limit_order).transmit = True
        else:
            # Only main order
            main_order.transmit = True

        # Store orders
        orders = [order for order in [main_order, stop_order, limit_order] if order is not None]
        if not hasattr(self.backtrader_strategy, 'orders'):
            self.backtrader_strategy.orders = []
        self.backtrader_strategy.orders.extend(orders)

        # Update balances (simplified for backtesting)
        trade_amount = size * (order_price if order_price else price)
        if trade_type.lower() == 'buy':
            self.available_strategy_balance -= trade_amount
            self.available_balance -= trade_amount
            self.current_balance -= trade_amount
        else:
            self.available_strategy_balance += trade_amount
            self.available_balance += trade_amount
            self.current_balance += trade_amount

        message = f"{action} order executed for {size} {symbol} at {order_type} price."
        self.notify_user(message)
        logger.info(message)

        # Update exec_context and save context
        self.exec_context['available_strategy_balance'] = self.available_strategy_balance
        self.exec_context['available_balance'] = self.available_balance
        self.exec_context['current_balance'] = self.current_balance
        self.save_context()

    # 2. Override process_indicator
    def process_indicator(self, indicator_name: str, output_field: str):
        """
        Retrieves precomputed indicator values for backtesting.
        If the current value is NaN, returns the last non-NaN value if available.
        If no last valid value exists, searches forward for the next valid value.
        If no valid value is found, returns a default value (e.g., 1).
        """
        logger.debug(f"Backtester is retrieving indicator '{indicator_name}' from precomputed data.")
        if self.backtrader_strategy is None:
            logger.error("Backtrader strategy is not set in StrategyInstance.")
            return None

        df = self.backtrader_strategy.precomputed_indicators.get(indicator_name)
        if df is None:
            logger.error(f"Indicator '{indicator_name}' not found.")
            return None

        idx = self.backtrader_strategy.indicator_pointers.get(indicator_name, 0)
        if idx >= len(df):
            logger.warning(f"No more data for indicator '{indicator_name}' at index {idx}.")
            return None

        # Retrieve the value at the current index
        value = df.iloc[idx].get(output_field)

        if pd.isna(value):
            # Check if we have a cached last valid value
            last_valid_value = self.last_valid_values.get(indicator_name, {}).get(output_field)
            if last_valid_value is not None:
                logger.debug(f"Using cached last valid value for indicator '{indicator_name}': {last_valid_value}")
                return last_valid_value
            else:
                logger.debug(
                    f"No cached last valid value for indicator '{indicator_name}'. Searching ahead for next valid value.")
                # Search forward for the next valid value
                valid_idx = idx + 1
                while valid_idx < len(df):
                    next_value = df.iloc[valid_idx].get(output_field)
                    if not pd.isna(next_value):
                        logger.debug(f"Found valid value at index {valid_idx}: {next_value}")
                        # Update the cache with this value
                        if indicator_name not in self.last_valid_values:
                            self.last_valid_values[indicator_name] = {}
                        self.last_valid_values[indicator_name][output_field] = next_value
                        return next_value
                    valid_idx += 1
                # If no valid value is found, return a default value (e.g., 1)
                logger.warning(
                    f"No valid value found for indicator '{indicator_name}' after index {idx}. Returning default value 1.")
                return 1  # Default value to prevent errors
        else:
            # Update the cache with the new valid value
            if indicator_name not in self.last_valid_values:
                self.last_valid_values[indicator_name] = {}
            self.last_valid_values[indicator_name][output_field] = value
            return value

    # 3. Override get_current_price
    def get_current_price(self, timeframe: str = '1h', exchange: str = 'binance',
                          symbol: str = 'BTC/USD') -> float:
        """
        Retrieves the current market price from Backtrader's data feed.
        """
        if self.backtrader_strategy:
            price = self.backtrader_strategy.data.close[0]
            logger.debug(f"Current price from Backtrader's data feed: {price}")
            return price
        else:
            logger.error("Backtrader strategy is not set.")
            return 0.0

    # 4. Override get_last_candle
    def get_last_candle(self, candle_part: str, timeframe: str, exchange: str, symbol: str):
        """
        Retrieves the specified part of the last candle from Backtrader's data feed.
        """
        if self.backtrader_strategy is None:
            logger.error("Backtrader strategy is not set in StrategyInstance.")
            return None

        candle_map = {
            'open': self.backtrader_strategy.data.open[0],
            'high': self.backtrader_strategy.data.high[0],
            'low': self.backtrader_strategy.data.low[0],
            'close': self.backtrader_strategy.data.close[0],
            'volume': self.backtrader_strategy.data.volume[0],
        }
        value = candle_map.get(candle_part.lower())
        if value is None:
            logger.error(f"Invalid candle_part '{candle_part}'. Must be one of {list(candle_map.keys())}.")
        else:
            logger.debug(
                f"Retrieved '{candle_part}' from last candle for {symbol} on {exchange} ({timeframe}): {value}"
            )
        return value

    # 5. Override get_filled_orders
    def get_filled_orders(self) -> int:
        """
        Retrieves the number of filled orders from Backtrader's broker.
        """
        if self.backtrader_strategy is None:
            logger.error("Backtrader strategy is not set in StrategyInstance.")
            return 0
        try:
            filled_orders = len([o for o in self.backtrader_strategy.broker.orders if o.status == bt.Order.Completed])
            logger.debug(f"Number of filled orders: {filled_orders}")
            return filled_orders
        except Exception as e:
            logger.error(f"Error retrieving filled orders: {e}", exc_info=True)
            return 0

    # 6. Override get_available_balance
    def get_available_balance(self) -> float:
        """
        Retrieves the available cash balance from Backtrader's broker.
        """
        self.available_balance = self.broker.getcash()
        self.exec_context['available_balance'] = self.available_balance
        logger.debug(f"Available balance retrieved from Backtrader's broker: {self.available_balance}")
        return self.available_balance

    # 7. Override get_current_balance
    def get_current_balance(self) -> float:
        """
        Retrieves the current total value from Backtrader's broker.
        """
        self.current_balance = self.broker.getvalue()
        self.exec_context['current_balance'] = self.current_balance
        logger.debug(f"Current balance retrieved from Backtrader's broker: {self.current_balance}")
        return self.current_balance

    # 8. Override get_filled_orders_details (Optional but Recommended)
    def get_filled_orders_details(self) -> list:
        """
        Retrieves detailed information about filled orders.
        """
        if self.backtrader_strategy is None:
            logger.error("Backtrader strategy is not set in StrategyInstance.")
            return []
        try:
            filled_orders = []
            for order in self.backtrader_strategy.broker.filled:
                order_info = {
                    'ref': order.ref,
                    'size': order.size,
                    'price': order.executed.price,
                    'value': order.executed.value,
                    'commission': order.executed.comm,
                    'status': order.status,
                    'created_at': dt.datetime.fromtimestamp(order.created.dt.timestamp()) if hasattr(order,
                                                                                                     'created') else None
                }
                filled_orders.append(order_info)
            logger.debug(f"Filled orders details: {filled_orders}")
            return filled_orders
        except Exception as e:
            logger.error(f"Error retrieving filled orders details: {e}", exc_info=True)
            return []

    # 9. Override notify_user
    def notify_user(self, message: str):
        """
        Suppresses user notifications and instead logs them.
        :param message: Notification message.
        """
        logger.debug(f"Backtest notification: {message}")

    def save_context(self):
        """
        Saves the current strategy execution context to the cache and database.
        Adjusted for backtesting to include balance attributes.
        """
        try:
            # Update balances from broker before saving
            self.current_balance = self.get_current_balance()
            self.available_balance = self.get_available_balance()

            self.exec_context['current_balance'] = self.current_balance
            self.exec_context['available_balance'] = self.available_balance
            self.exec_context['available_strategy_balance'] = self.available_strategy_balance

            super().save_context()
        except Exception as e:
            logger.error(f"Error saving context for backtest: {e}", exc_info=True)

    def fetch_user_balance(self) -> float:
        """
        Fetches the starting balance from Backtrader's broker.
        """
        if hasattr(self, 'broker') and self.broker:
            balance = self.broker.getvalue()
            logger.debug(f"Fetched starting balance from Backtrader's broker: {balance}")
            return balance
        else:
            logger.error("Broker is not set. Cannot fetch starting balance.")
            return 0.0


    def calculate_available_balance(self) -> float:
        """
        Calculates the available cash balance from Backtrader's broker.
        """
        if self.broker:
            available_balance = self.broker.getcash()
            logger.debug(f"Calculated available cash balance from Backtrader's broker: {available_balance}")
            return available_balance
        else:
            logger.error("Broker is not set. Cannot calculate available balance.")
            return 0.0


    def set_available_strategy_balance(self, balance: float):
        """
        Sets the available strategy balance in backtesting.
        """
        # In backtesting, we might simulate allocation by adjusting internal variables
        if balance > self.get_available_balance():
            raise ValueError("Cannot allocate more than the available balance in backtest.")
        self.available_strategy_balance = balance
        self.exec_context['available_strategy_balance'] = self.available_strategy_balance
        self.save_context()
        logger.debug(f"Available strategy balance set to {balance} in backtest.")

    def get_available_strategy_balance(self) -> float:
        """
        Retrieves the available strategy balance in backtesting.
        """
        logger.debug(f"Available strategy balance in backtest: {self.available_strategy_balance}")
        return self.available_strategy_balance

    def get_starting_balance(self) -> float:
        """
        Returns the starting balance in backtesting.
        """
        logger.debug(f"Starting balance in backtest: {self.starting_balance}")
        return self.starting_balance

    def get_active_trades(self) -> int:
        """
        Retrieves the number of active trades (open positions) from Backtrader's broker.
        """
        if self.backtrader_strategy is None:
            logger.error("Backtrader strategy is not set in StrategyInstance.")
            return 0
        try:
            # Get all positions
            positions = self.broker.positions
            active_trades_count = sum(1 for position in positions.values() if position.size != 0)
            logger.debug(f"Number of active trades: {active_trades_count}")
            return active_trades_count
        except Exception as e:
            logger.error(f"Error retrieving active trades: {e}", exc_info=True)
            return 0
