# backtesting.py

import logging
import math
import time
import uuid
import eventlet

import backtrader as bt
import datetime as dt
from DataCache_v3 import DataCache, RowBasedCache, TableBasedCache
from Strategies import Strategies
from StrategyInstance import StrategyInstance
from indicators import Indicators
from backtest_strategy_instance import BacktestStrategyInstance
from mapped_strategy import MappedStrategy
import numpy as np
import pandas as pd
import signal
from utils import sanitize_for_json

# Configure logging
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)  # Set to DEBUG for detailed logging
handler = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)


# Custom EquityCurveAnalyzer
class EquityCurveAnalyzer(bt.Analyzer):
    def __init__(self):
        self.equity_curve = []

    def next(self):
        self.equity_curve.append(self.strategy.broker.getvalue())

    def get_analysis(self):
        return {'equity_curve': self.equity_curve}

# Backtester Class
class Backtester:
    def __init__(self, data_cache: DataCache, strategies: Strategies, indicators: Indicators, socketio,
                 edm_client=None, external_indicators=None, signals=None, formations=None):
        """ Initialize the Backtesting class with a cache for back-tests """
        self.data_cache = data_cache
        self.strategies = strategies
        self.indicators_manager = indicators
        self.socketio = socketio
        self.edm_client = edm_client
        self.external_indicators = external_indicators
        self.signals = signals
        self.formations = formations

        # Ensure 'tests' cache exists
        self.data_cache.create_cache(
            'tests',
            cache_type='row',
            size_limit=100,
            default_expiration=dt.timedelta(hours=2),  # Set to 2 hours
            eviction_policy='evict',
            columns=[
                'user_name',  # TEXT
                'strategy_name',  # TEXT
                'start_time',  # TEXT or DATETIME
                'capital',  # REAL or INTEGER
                'commission',  # REAL
                'results',  # TEXT or JSON
                'strategy_instance_id'  # TEXT
            ]
        )
        # Start periodic purge (optional)
        self.start_periodic_purge(interval_seconds=3600)  # Purge every hour

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

    @staticmethod
    def _coerce_user_id(user_id) -> int | None:
        if user_id is None or user_id == '':
            return None
        try:
            return int(user_id)
        except (TypeError, ValueError):
            return None

    def _resolve_user_name(self, msg_data: dict) -> str | None:
        user_name = msg_data.get('user_name') or msg_data.get('user')
        if user_name:
            return user_name

        user_id = self._coerce_user_id(msg_data.get('user_id') or msg_data.get('userId'))
        if user_id is None:
            return None

        try:
            return self.data_cache.get_datacache_item(
                item_name='user_name',
                cache_name='users',
                filter_vals=('id', user_id)
            )
        except Exception:
            logger.warning(f"Unable to resolve user_name from user id '{user_id}'.")
            return None

    def _resolve_user_id(self, msg_data: dict, user_id=None, user_name: str | None = None) -> int | None:
        normalized_user_id = self._coerce_user_id(user_id)
        if normalized_user_id is not None:
            return normalized_user_id

        normalized_user_id = self._coerce_user_id(msg_data.get('user_id') or msg_data.get('userId'))
        if normalized_user_id is not None:
            return normalized_user_id

        if user_name:
            try:
                return self.data_cache.get_datacache_item(
                    item_name='id',
                    cache_name='users',
                    filter_vals=('user_name', user_name)
                )
            except Exception:
                logger.warning(f"Unable to resolve user_id from user_name '{user_name}'.")
                return None
        return None

    def cache_backtest(self, backtest_key: str, backtest_data: dict, strategy_instance_id: str):
        """
        Cache the backtest data for a user.
        If the backtest already exists, update it; otherwise, insert a new entry.
        :param backtest_key: A unique identifier for the backtest.
        :param backtest_data: Dictionary containing backtest parameters.
        :param strategy_instance_id: The ID of the strategy instance.
        """
        # Define columns and corresponding values (excluding 'tbl_key')
        columns = ('user_name', 'strategy_name', 'start_time', 'capital', 'commission', 'results', 'strategy_instance_id')
        values = (
            backtest_data.get('user_name'),
            backtest_data.get('strategy'),
            backtest_data.get('start_date'),
            backtest_data.get('capital', 10000),      # Default capital if not provided
            backtest_data.get('commission', 0.001),  # Default commission
            None,  # No results yet; will be filled in after backtest completion
            strategy_instance_id
        )

        try:
            # Check if the backtest already exists
            existing_backtest = self.data_cache.get_rows_from_cache(cache_name='tests',
                                                                    filter_vals=[('tbl_key', backtest_key)])

            if existing_backtest.empty:
                # Insert new backtest entry
                self.data_cache.insert_row_into_cache(
                    cache_name='tests', columns=columns, values=values, key=backtest_key)
                logger.debug(f"Inserted new backtest entry '{backtest_key}'.")
            else:
                # Update existing backtest entry (e.g., reset 'results' if needed)
                # Here, we assume you might want to reset 'results' when re-running
                self.data_cache.modify_datacache_item(
                    cache_name='tests',
                    filter_vals=[('tbl_key', backtest_key)],
                    field_names=('results',),
                    new_values=(None,),  # Reset results
                    overwrite='tbl_key'  # Ensures uniqueness based on 'tbl_key'
                )
                logger.debug(f"Updated existing backtest entry '{backtest_key}'. Reset 'results'.")
        except Exception as e:
            logger.error(f"Error in cache_backtest for '{backtest_key}': {e}", exc_info=True)


    def cleanup_backtest(self, backtest_key: str, strategy_instance_id: str) -> None:
        """
        Removes cache entries related to a specific backtest and its strategy context.

        :param backtest_key: The unique key identifying the backtest in the 'tests' cache.
        :param strategy_instance_id: The unique identifier for the strategy instance in 'strategy_contexts'.
        """
        try:
            # Remove backtest entry from 'tests' cache
            self.data_cache.remove_row_from_datacache(
                cache_name='tests',
                filter_vals=[('tbl_key', backtest_key)],
                remove_from_db=False  # Assuming tests are transient and not stored in DB
            )
            logger.info(f"Backtest '{backtest_key}' removed from 'tests' cache.")

            # Remove strategy context from 'strategy_contexts' cache
            # Assuming 'strategy_contexts' is a table-based cache
            self.data_cache.remove_row_from_datacache(
                cache_name='strategy_contexts',
                filter_vals=[('strategy_instance_id', strategy_instance_id)],
                remove_from_db=True  # Remove from DB as well
            )
            logger.info(f"Strategy context '{strategy_instance_id}' removed from 'strategy_contexts' cache.")
        except Exception as e:
            logger.error(f"Error during cleanup of backtest '{backtest_key}': {e}", exc_info=True)

    def remove_backtest(self, backtest_key: str):
        """
        Remove a backtest from the 'tests' cache.
        :param backtest_key: The unique key identifying the backtest.
        """
        try:
            self.data_cache.remove_row_from_datacache(
                cache_name='tests',
                filter_vals=[('tbl_key', backtest_key)],
                remove_from_db = False  # Assuming tests are transient and not stored in DB
            )
            logger.info(f"Backtest '{backtest_key}' removed from 'tests' cache.")
        except KeyError:
            logger.error(f"Backtest '{backtest_key}' not found in 'tests' cache.")
            raise KeyError(f"Backtest '{backtest_key}' not found.")
        except Exception as e:
            logger.error(f"Error removing backtest '{backtest_key}': {e}", exc_info=True)
            raise

    def validate_strategy_components(self, tbl_key: str, user_name: str) -> dict:
        """
        Retrieves and validates the components of a user-defined strategy.
        Raises a ValueError if required components are missing or incorrectly formatted.

        :param tbl_key: The unique identifier of the strategy.
        :param user_name: The username associated with the strategy.
        :return: The validated strategy data.
        """
        try:
            user_strategy = self.strategies.get_strategy_by_tbl_key(tbl_key)
        except ValueError:
            logger.error(f"Invalid tbl_key '{tbl_key}'.")
            raise ValueError(f"Invalid tbl_key '{tbl_key}'.")

        if not user_strategy:
            logger.error(f"Strategy with tbl_key '{tbl_key}' not found for user '{user_name}'.")
            raise ValueError(f"Strategy with tbl_key '{tbl_key}' not found for user '{user_name}'.")

        strategy_components = user_strategy.get('strategy_components', {})
        generated_code = strategy_components.get('generated_code')
        indicators_used = strategy_components.get('indicators')

        # Validate extracted data
        if not generated_code:
            logger.error("No 'generated_code' found in strategy components.")
            raise ValueError("Strategy must contain 'generated_code'.")

        if not isinstance(indicators_used, list):
            logger.error("'indicators_used' should be a list.")
            raise ValueError("'indicators_used' should be a list.")

        return user_strategy


    def prepare_data_feed(self, start_date: str, source, user_name: str) -> pd.DataFrame:
        """
        Prepare the main data feed based on the start date and source.

        :param start_date: Start date in 'YYYY-MM-DDTHH:MM' format.
        :param source: Can be either a list or a dictionary.
                       - If a list: Expected order is [exchange, symbol, timeframe].
                       - If a dictionary: Expected keys are 'exchange', 'symbol', and 'timeframe'.
        :param user_name: The user name associated with the data feed.
        :return: Pandas DataFrame with OHLC data.

        This method is designed to be flexible, supporting both list and dictionary formats
        for the source. This flexibility allows for backward compatibility with existing code
        using lists, while providing a clearer structure when dictionaries are preferred.
        """
        try:
            # Check if start_date is present and non-empty
            if not start_date:
                # Set default start_date to one hour ago in UTC
                one_hour_ago = dt.datetime.now(dt.timezone.utc) - dt.timedelta(hours=1)
                start_date = one_hour_ago.strftime('%Y-%m-%dT%H:%M')
                logger.warning(f"No 'start_date' provided. Setting default start_date to '{start_date}'.")

            # Convert the start date to a datetime object and make it timezone-aware (UTC)
            start_dt = dt.datetime.strptime(start_date, '%Y-%m-%dT%H:%M')
            start_dt = start_dt.replace(tzinfo=dt.timezone.utc)  # Set UTC timezone

            # Check if source is a dictionary or a list, then set exchange, symbol, and timeframe
            if isinstance(source, dict):
                exchange = source.get('exchange', 'binance')
                symbol = source.get('symbol', 'BTC/USDT')
                timeframe = source.get('timeframe', '1h')
            elif isinstance(source, list) and len(source) >= 3:
                exchange, symbol, timeframe = source[0], source[1], source[2]
            else:
                logger.error("Source must be either a list with at least 3 elements or a dictionary.")
                return pd.DataFrame()

            # Fetch OHLC data from EDM (Exchange Data Manager)
            if self.edm_client is None:
                raise RuntimeError("EDM client not initialized. Cannot fetch backtest data.")

            logger.info(f"Fetching backtest data from EDM: {exchange}/{symbol}/{timeframe} from {start_dt}")

            try:
                data = self.edm_client.get_candles_sync(
                    exchange=exchange.lower(),
                    symbol=symbol,
                    timeframe=timeframe,
                    start=start_dt,
                    end=dt.datetime.now(dt.timezone.utc)
                )
            except Exception as e:
                logger.error(f"EDM get_candles_sync raised exception: {e}", exc_info=True)
                return pd.DataFrame()

            logger.info(f"EDM returned {len(data)} rows for {exchange}/{symbol}/{timeframe}")

            if data.empty:
                logger.error(
                    f"No data retrieved for symbol {symbol} on exchange {exchange} with timeframe {timeframe}.")
                return pd.DataFrame()  # Return empty DataFrame

            logger.info(f"Data feed prepared for {symbol} on {exchange} with timeframe {timeframe}. "
                       f"Retrieved {len(data)} candles.")

            return data

        except Exception as e:
            logger.error(f"Error preparing data feed: {e}")
            return pd.DataFrame()

    def precompute_indicators(self, indicators_definitions: list, user_name: str, data_feed: pd.DataFrame,
                               indicator_owner_id: int = None) -> dict:
        """
        Precompute indicator values directly on the backtest data feed.

        IMPORTANT: This computes indicators on the actual backtest candle data,
        ensuring the indicator values align with the price data used in the backtest.
        Previously, this fetched fresh/latest candles which caused misalignment.

        :param indicators_definitions: List of indicator definitions needed.
        :param user_name: The user running the backtest.
        :param data_feed: The candle data for backtesting.
        :param indicator_owner_id: For subscribed strategies, the creator's user ID for indicator lookup.
        """
        import json as json_module  # Local import to avoid conflicts

        precomputed_indicators = {}
        total_candles = len(data_feed)

        logger.info(f"[BACKTEST] precompute_indicators called with {len(indicators_definitions)} indicator definitions")
        logger.info(f"[BACKTEST] user_name: {user_name}, total_candles: {total_candles}")
        logger.info(f"[BACKTEST] data_feed columns: {list(data_feed.columns)}")
        logger.info(f"[BACKTEST] data_feed first row: {data_feed.iloc[0].to_dict() if len(data_feed) > 0 else 'empty'}")
        logger.info(f"[BACKTEST] data_feed last row: {data_feed.iloc[-1].to_dict() if len(data_feed) > 0 else 'empty'}")

        # Aggregate requested outputs for each indicator
        indicator_outputs = {}
        for indicator_def in indicators_definitions:
            indicator_name = indicator_def.get('name')
            output = indicator_def.get('output')

            if not indicator_name:
                logger.warning("Indicator definition missing 'name'. Skipping.")
                continue

            # Initialize the outputs set if necessary
            if indicator_name not in indicator_outputs:
                indicator_outputs[indicator_name] = set()

            if output:
                # If outputs is not None, add the output to the set
                if indicator_outputs[indicator_name] is not None:
                    indicator_outputs[indicator_name].add(output)
            else:
                # If output is None, we need all outputs
                indicator_outputs[indicator_name] = None  # None indicates all outputs

        # Get user ID for indicator lookup
        # For subscribed strategies, use indicator_owner_id (creator's ID) instead of running user's ID
        if indicator_owner_id is not None:
            user_id = indicator_owner_id
            logger.info(f"[BACKTEST] Using indicator_owner_id {indicator_owner_id} for indicator lookup (subscribed strategy)")
        else:
            user_id = self.data_cache.get_datacache_item(
                item_name='id',
                cache_name='users',
                filter_vals=('user_name', user_name)
            )

        logger.info(f"[BACKTEST] indicator_outputs to precompute: {indicator_outputs}")

        # Prepare candle data for indicator calculation
        # Convert data_feed to the format expected by indicators
        candle_data = data_feed.copy()

        # Ensure required columns exist
        if 'time' not in candle_data.columns and candle_data.index.name == 'datetime':
            candle_data['time'] = candle_data.index.astype(np.int64) // 10**6  # Convert to milliseconds

        for indicator_name, outputs in indicator_outputs.items():
            logger.info(f"[BACKTEST] Computing indicator '{indicator_name}' on backtest data feed")

            try:
                # First, check if this is an external indicator (historical API)
                if self.external_indicators:
                    ext_indicator = self.external_indicators.get_indicator_by_name(indicator_name, user_id)
                    if ext_indicator:
                        logger.info(f"[BACKTEST] '{indicator_name}' is an external indicator - fetching historical data from API")
                        # Use the external indicator's calculate method
                        result_series = ext_indicator.calculate(candle_data)

                        if result_series is not None and len(result_series) > 0:
                            # Convert series to DataFrame with 'value' column
                            result_df = pd.DataFrame({
                                'time': candle_data['time'] if 'time' in candle_data.columns else range(len(result_series)),
                                'value': result_series.values
                            })
                            result_df.reset_index(drop=True, inplace=True)
                            precomputed_indicators[indicator_name] = result_df
                            logger.info(f"[BACKTEST] Precomputed external indicator '{indicator_name}' with {len(result_df)} data points.")
                        else:
                            logger.warning(f"[BACKTEST] No historical data available for external indicator '{indicator_name}'.")
                        continue

                # Fetch indicator definition from cache
                indicators = self.data_cache.get_rows_from_datacache(
                    cache_name='indicators',
                    filter_vals=[('creator', str(user_id)), ('name', indicator_name)]
                )

                if indicators.empty:
                    logger.warning(f"[BACKTEST] Indicator '{indicator_name}' not found for user '{user_name}' (id={user_id}). Skipping.")
                    continue

                indicator = indicators.iloc[0]
                kind = indicator['kind']
                properties = json_module.loads(indicator['properties']) if isinstance(indicator['properties'], str) else indicator['properties']

                logger.info(f"[BACKTEST] Indicator '{indicator_name}' is of kind '{kind}' with properties: {properties}")

                # Get the indicator class from the registry
                indicator_class = self.indicators_manager.indicator_registry.get(kind)
                if not indicator_class:
                    logger.warning(f"[BACKTEST] Unknown indicator kind '{kind}' for '{indicator_name}'. Skipping.")
                    continue

                # Instantiate and calculate
                indicator_obj = indicator_class(name=indicator_name, indicator_type=kind, properties=properties)
                result_df = indicator_obj.calculate(candles=candle_data, user_name=user_name, num_results=total_candles)

                if result_df is None or (isinstance(result_df, pd.DataFrame) and result_df.empty):
                    logger.warning(f"[BACKTEST] No data computed for indicator '{indicator_name}'. Skipping.")
                    continue

                logger.info(f"[BACKTEST] Computed indicator '{indicator_name}': {len(result_df)} rows, columns: {list(result_df.columns)}")

                # Log first few values for debugging
                if len(result_df) > 0:
                    logger.info(f"[BACKTEST] First 3 rows of '{indicator_name}': {result_df.head(3).to_dict('records')}")
                    logger.info(f"[BACKTEST] Last 3 rows of '{indicator_name}': {result_df.tail(3).to_dict('records')}")

                # If outputs is None, keep all outputs
                if outputs is not None:
                    # Include 'time' and requested outputs
                    columns_to_keep = ['time'] + list(outputs)
                    missing_columns = [col for col in columns_to_keep if col not in result_df.columns]
                    if missing_columns:
                        logger.warning(f"[BACKTEST] Indicator '{indicator_name}' missing columns: {missing_columns}. Available: {list(result_df.columns)}")
                        # Try to continue with available columns
                        columns_to_keep = [c for c in columns_to_keep if c in result_df.columns]
                    result_df = result_df[columns_to_keep]

                # Reset index and store the DataFrame
                result_df.reset_index(drop=True, inplace=True)
                precomputed_indicators[indicator_name] = result_df
                logger.info(f"[BACKTEST] Precomputed indicator '{indicator_name}' with {len(result_df)} data points.")

            except Exception as e:
                logger.error(f"[BACKTEST] Error computing indicator '{indicator_name}': {e}", exc_info=True)
                continue

        return precomputed_indicators

    def _calculate_warmup_period(self, indicators_definitions: list, user_name: str, indicator_owner_id: int = None) -> int:
        """
        Calculate the maximum warmup period needed based on indicator periods.

        :param indicators_definitions: List of indicator definitions from strategy
        :param user_name: Username for looking up indicator configs
        :param indicator_owner_id: For subscribed strategies, the creator's user ID for indicator lookup.
        :return: Maximum warmup period in candles
        """
        import json as json_module

        max_period = 0

        # For subscribed strategies, use indicator_owner_id (creator's ID) instead of running user's ID
        if indicator_owner_id is not None:
            user_id = indicator_owner_id
            logger.info(f"[BACKTEST] Using indicator_owner_id {indicator_owner_id} for warmup calculation (subscribed strategy)")
        else:
            user_id = self.data_cache.get_datacache_item(
                item_name='id',
                cache_name='users',
                filter_vals=('user_name', user_name)
            )

        for indicator_def in indicators_definitions:
            indicator_name = indicator_def.get('name')
            if not indicator_name:
                continue

            try:
                indicators = self.data_cache.get_rows_from_datacache(
                    cache_name='indicators',
                    filter_vals=[('creator', str(user_id)), ('name', indicator_name)]
                )

                if indicators.empty:
                    continue

                indicator = indicators.iloc[0]
                properties = json_module.loads(indicator['properties']) if isinstance(indicator['properties'], str) else indicator['properties']

                # Get period from properties (most indicators use 'period')
                period = properties.get('period', 0)
                if period > max_period:
                    max_period = period

            except Exception as e:
                logger.warning(f"Could not get period for indicator '{indicator_name}': {e}")
                continue

        logger.info(f"[BACKTEST] Maximum indicator warmup period: {max_period}")
        return max_period

    def _get_timeframe_minutes(self, timeframe: str) -> int:
        """Convert timeframe string to minutes."""
        timeframe_map = {
            '1m': 1, '3m': 3, '5m': 5, '15m': 15, '30m': 30,
            '1h': 60, '2h': 120, '4h': 240, '6h': 360, '12h': 720,
            '1d': 1440, '1w': 10080
        }
        return timeframe_map.get(timeframe.lower(), 60)  # Default to 1h

    def prepare_backtest_data(self, msg_data: dict, strategy_components: dict, indicator_owner_id: int = None) -> tuple:
        """
        Prepare the data feed and precomputed indicators for backtesting.
        :param msg_data: Message data containing backtest parameters.
        :param strategy_components: Components of the user-defined strategy.
        :param indicator_owner_id: For subscribed strategies, the creator's user ID for indicator lookup.
        :return: Tuple of (data_feed, precomputed_indicators).
        :raises ValueError: If data sources are invalid or data feed cannot be prepared.
        """
        user_name = self._resolve_user_name(msg_data)
        if not user_name:
            raise ValueError("User not specified in backtest request.")

        data_sources = strategy_components.get('data_sources', [])

        logger.info(f"Strategy data_sources: {data_sources}")

        if not data_sources:
            logger.error("No valid data sources found in the strategy.")
            raise ValueError("No valid data sources found in the strategy.")

        # For simplicity, use the first data source as the main data feed.
        main_source = data_sources[0]
        logger.info(f"Using main_source for backtest: {main_source}")

        # Calculate warmup period needed for indicators
        indicators_definitions = strategy_components.get('indicators', [])
        warmup_candles = self._calculate_warmup_period(indicators_definitions, user_name, indicator_owner_id=indicator_owner_id)

        # Get timeframe to calculate how far back to fetch for warmup
        timeframe = main_source.get('timeframe', '1h')
        timeframe_minutes = self._get_timeframe_minutes(timeframe)
        warmup_minutes = warmup_candles * timeframe_minutes

        # Prepare the main data feed with extended start for warmup
        original_start_date = msg_data.get('start_date', '2023-01-01T00:00')
        logger.info(f"Backtest original start_date: {original_start_date}")

        # Calculate adjusted start date for indicator warmup
        original_start_dt = dt.datetime.strptime(original_start_date, '%Y-%m-%dT%H:%M')
        adjusted_start_dt = original_start_dt - dt.timedelta(minutes=warmup_minutes)
        adjusted_start_date = adjusted_start_dt.strftime('%Y-%m-%dT%H:%M')

        logger.info(f"[BACKTEST] Fetching data from {adjusted_start_date} (adjusted for {warmup_candles} warmup candles)")

        data_feed = self.prepare_data_feed(adjusted_start_date, main_source, user_name)

        if data_feed.empty:
            logger.error("Data feed could not be prepared. Please check the data source.")
            raise ValueError("Data feed could not be prepared. Please check the data source.")

        # Precompute indicator values on the full dataset (including warmup candles)
        precomputed_indicators = self.precompute_indicators(
            indicators_definitions, user_name, data_feed, indicator_owner_id=indicator_owner_id
        )

        # Now trim BOTH the data feed AND indicators to start at the user's original start_date
        # This ensures the first indicator values in the backtest have full warmup context
        if precomputed_indicators:
            # Find where the original start_date falls in the data
            original_start_unix = original_start_dt.replace(tzinfo=dt.timezone.utc).timestamp()

            # Find the index where we should start the backtest
            backtest_start_idx = 0
            for idx, row in data_feed.iterrows():
                if row['time'] >= original_start_unix:
                    backtest_start_idx = idx
                    break

            # Calculate how many indicator rows to skip
            min_indicator_len = min(len(df) for df in precomputed_indicators.values())
            original_feed_len = len(data_feed)
            indicator_warmup = original_feed_len - min_indicator_len

            # The effective backtest start is max of (user's start, indicator warmup)
            effective_start_idx = max(backtest_start_idx, indicator_warmup)

            logger.info(f"[BACKTEST] Original data length: {original_feed_len}, indicator warmup: {indicator_warmup}, user start idx: {backtest_start_idx}, effective start: {effective_start_idx}")

            # Trim data feed
            if effective_start_idx > 0:
                data_feed = data_feed.iloc[effective_start_idx:].reset_index(drop=True)
                logger.info(f"[BACKTEST] Trimmed data feed to {len(data_feed)} rows starting from effective start")

            # Trim indicators to match
            # Indicators already have warmup_period fewer rows, so we need to adjust
            indicator_trim = effective_start_idx - indicator_warmup
            if indicator_trim > 0:
                for name, df in precomputed_indicators.items():
                    precomputed_indicators[name] = df.iloc[indicator_trim:].reset_index(drop=True)
                    logger.info(f"[BACKTEST] Trimmed indicator '{name}' to {len(precomputed_indicators[name])} rows")

        logger.info("Backtest data prepared successfully.")

        return data_feed, precomputed_indicators

    def run_backtest(self, strategy_class, data_feed: pd.DataFrame, msg_data: dict, user_name: str,
                     socket_conn_id: str, strategy_instance: BacktestStrategyInstance, backtest_name: str,
                     user_id: str, backtest_key: str, strategy_name: str, precomputed_indicators: dict,
                     precomputed_signals: dict = None):
        """
        Runs a backtest using Backtrader and uses Flask-SocketIO's background tasks.
        Sends progress updates to the client via WebSocket.
        """

        tbl_key = msg_data.get('strategy')  # Expecting tbl_key instead of strategy_name

        def execute_backtest():
            try:
                # **Convert 'time' to 'datetime' if necessary**
                if 'time' in data_feed.columns:
                    # Time values from EDM are Unix timestamps in SECONDS, not milliseconds
                    data_feed['datetime'] = pd.to_datetime(data_feed['time'], unit='s')
                    # DEBUG: Log first and last timestamps to verify conversion
                    if len(data_feed) > 0:
                        first_time = data_feed['time'].iloc[0]
                        last_time = data_feed['time'].iloc[-1]
                        first_dt = data_feed['datetime'].iloc[0]
                        last_dt = data_feed['datetime'].iloc[-1]
                        logger.info(f"[DEBUG DATETIME FIX] First raw time: {first_time}, converted: {first_dt}")
                        logger.info(f"[DEBUG DATETIME FIX] Last raw time: {last_time}, converted: {last_dt}")
                    data_feed.set_index('datetime', inplace=True)
                    logger.info("Converted 'time' to 'datetime' and set as index in data_feed.")

                # **Select relevant columns for Backtrader**
                columns_to_keep = ['open', 'high', 'low', 'close', 'volume']
                if not set(columns_to_keep).issubset(data_feed.columns):
                    logger.error("Data feed is missing one or more required columns: %s", columns_to_keep)
                    raise ValueError("Incomplete data feed for Backtrader.")

                cerebro = bt.Cerebro()

                # Assign cerebro to strategy_instance for potential use in custom methods
                strategy_instance.cerebro = cerebro

                # Add the mapped strategy to the backtest, including strategy_instance and precomputed_indicators as parameters
                cerebro.addstrategy(
                    strategy_class,
                    strategy_instance=strategy_instance,
                    precomputed_indicators=precomputed_indicators,  # Pass precomputed indicators
                    precomputed_signals=precomputed_signals or {},  # Pass precomputed signals
                    socketio=self.socketio,  # Pass SocketIO instance
                    socket_conn_id=socket_conn_id,  # Pass SocketIO connection ID
                    data_length=len(data_feed),  # Pass data length for progress updates
                    backtest_name=backtest_name  # Pass backtest_name for progress updates
                )

                # Add data feed to Cerebro
                bt_feed = bt.feeds.PandasData(dataname=data_feed)
                cerebro.adddata(bt_feed)

                # Set initial capital and commission
                initial_cash = msg_data.get('capital', 10000)
                cerebro.broker.setcash(initial_cash)
                commission = msg_data.get('commission', 0.001)
                cerebro.broker.setcommission(commission=commission)

                # Add analyzers
                cerebro.addanalyzer(EquityCurveAnalyzer, _name='equity_curve')
                cerebro.addanalyzer(bt.analyzers.TradeAnalyzer, _name='trade_analyzer')

                # Run the backtest
                logger.info("Running backtest...")
                start_time = dt.datetime.now()
                results = cerebro.run()
                end_time = dt.datetime.now()

                # Extract performance metrics
                final_value = cerebro.broker.getvalue()
                run_duration = (end_time - start_time).total_seconds()

                strategy = results[0]

                # Extract equity curve from analyzers
                equity_curve = strategy.analyzers.equity_curve.get_analysis().get('equity_curve', [])

                trades = strategy.trade_list

                # Send 100% completion
                self.socketio.emit('message', {'reply': 'progress',
                                               'data': { 'test_id': backtest_name,
                                                         'progress': 100}}
                                   , room=socket_conn_id)

                # Get collected alerts from strategy instance
                collected_alerts = strategy_instance.get_collected_alerts()

                # Get trading source info for chart validation
                trading_source = msg_data.get('trading_source', {})

                # Prepare the results to pass into the callback
                backtest_results = {
                    "success": True,  # Indicate success
                    "initial_capital": initial_cash,
                    "final_portfolio_value": final_value,
                    "run_duration": run_duration,
                    "equity_curve": equity_curve,
                    "trades": trades,
                    "alerts": collected_alerts,
                    "trading_source": trading_source,
                }

                logger.info("Backtest executed successfully.")

                # Invoke the callback with all necessary parameters
                self.backtest_callback(user_name, backtest_name, tbl_key,
                                       strategy_instance.strategy_instance_id, socket_conn_id,
                                       backtest_key, backtest_results)

            except Exception as e:
                # Handle exceptions and send error messages to the client
                error_message = f"Backtest execution failed: {str(e)}"
                self.socketio.emit('message', {'reply': 'backtest_error', 'data': {'message': error_message}},
                                   room=socket_conn_id)
                logger.error(f"[BACKTEST ERROR] {error_message}", exc_info=True)

                # Prepare failure results
                failure_results = {
                    "success": False,
                    "message": error_message
                }

                # Invoke callback with failure details to ensure cleanup
                self.backtest_callback(user_name, backtest_name, tbl_key,
                                       strategy_instance.strategy_instance_id, socket_conn_id,
                                       backtest_key, failure_results)

        # Start the backtest as a background task
        eventlet.spawn(execute_backtest)

    def handle_backtest_message(self, user_id: str, msg_data: dict, socket_conn_id: str) -> dict:
        """
        Handle incoming backtest messages, orchestrate the backtest process.
        """
        # Extract and define backtest parameters
        user_name = self._resolve_user_name(msg_data)
        user_id = self._resolve_user_id(msg_data, user_id=user_id, user_name=user_name)
        if not user_name or user_id is None:
            return {"error": "Missing user identity for backtest request."}

        msg_data.setdefault('user_name', user_name)
        msg_data.setdefault('user_id', user_id)
        msg_data.setdefault('userId', user_id)

        tbl_key = msg_data.get('strategy')  # Expecting tbl_key instead of strategy_name
        backtest_name = msg_data.get('backtest_name')  # Use the client-provided backtest_name

        # Authorization check: user must own the strategy OR be subscribed to it
        strategy = self.strategies.get_strategy_by_tbl_key(tbl_key)
        if not strategy:
            return {"error": "Strategy not found."}

        strategy_creator = strategy.get('creator')
        try:
            creator_id = int(strategy_creator) if strategy_creator is not None else None
        except (ValueError, TypeError):
            creator_id = None

        is_owner = (creator_id == user_id) if (creator_id is not None and user_id is not None) else False
        is_subscribed = self.strategies.is_subscribed(user_id, tbl_key)

        # Must be owner OR subscribed to run backtest
        if not is_owner and not is_subscribed:
            return {"error": "Subscribe to this strategy first"}

        # For subscribed strategies, use creator's indicators
        indicator_owner_id = creator_id if is_subscribed and not is_owner else None

        # Validate exchange requirements for backtest
        from exchange_validation import extract_required_exchanges, validate_for_backtest
        required_exchanges = extract_required_exchanges(strategy)

        if required_exchanges and self.edm_client:
            try:
                edm_available = self.edm_client.get_exchanges_sync()
                validation_result = validate_for_backtest(required_exchanges, edm_available)
                if not validation_result.valid:
                    return {
                        "error": validation_result.message,
                        "error_code": validation_result.error_code.value if validation_result.error_code else None,
                        "missing_exchanges": list(validation_result.missing_exchanges)
                    }
            except Exception as e:
                logger.warning(f"Could not validate EDM exchanges: {e}")
                return {
                    "error": "Cannot validate exchange availability - EDM unreachable",
                    "error_code": "edm_unreachable"
                }

        if not backtest_name:
            # If backtest_name is not provided, generate a unique name
            backtest_name = f"{tbl_key}_backtest"

        # To ensure uniqueness, append a UUID if necessary
        # Alternatively, rely on the client to provide a unique name
        # backtest_name = f"{backtest_name}_{uuid.uuid4()}"

        strategy_instance_id = f"test_{uuid.uuid4()}"
        backtest_key = f"backtest:{user_name}:{backtest_name}"

        # Retrieve the user strategy and validate it.
        try:
            user_strategy = self.validate_strategy_components(tbl_key, user_name)
        except ValueError as ve:
            return {"error": str(ve)}

        # Prepare the source and indicator feeds referenced in the strategy
        strategy_components = user_strategy.get('strategy_components', {})

        # Always use default_source as the primary data source if available.
        # This ensures we use the user's explicitly configured trading source,
        # even if data_sources was generated with incorrect symbol format.
        default_source = user_strategy.get('default_source')
        if default_source:
            # Map 'market' to 'symbol' if needed (default_source uses 'market', prepare_data_feed expects 'symbol')
            source = {
                'exchange': default_source.get('exchange'),
                'symbol': default_source.get('symbol') or default_source.get('market'),
                'timeframe': default_source.get('timeframe')
            }
            logger.info(f"Using default_source for backtest data: {source}")
            strategy_components['data_sources'] = [source]
            # Store trading source in msg_data for inclusion in backtest results
            msg_data['trading_source'] = source

        try:
            data_feed, precomputed_indicators = self.prepare_backtest_data(
                msg_data, strategy_components, indicator_owner_id=indicator_owner_id
            )
        except ValueError as ve:
            logger.error(f"Error preparing backtest data: {ve}")
            return {"error": str(ve)}

        # Instantiate BacktestStrategyInstance
        strategy_instance = BacktestStrategyInstance(
            strategy_instance_id=strategy_instance_id,
            strategy_id=tbl_key,  # Use tbl_key as strategy_id
            strategy_name=user_strategy.get("name"),
            user_id=user_id,
            generated_code=strategy_components.get("generated_code", ""),
            data_cache=self.data_cache,
            indicators=None,  # Custom handling in BacktestStrategyInstance
            trades=None,  # Custom handling in BacktestStrategyInstance
            indicator_owner_id=indicator_owner_id,  # For subscribed strategies, use creator's indicators
        )

        # Inject signals and formations for strategy execution
        strategy_instance.signals = self.signals
        strategy_instance.formations = self.formations
        strategy_instance.formation_owner_id = indicator_owner_id if indicator_owner_id else user_id

        # Cache the backtest
        self.cache_backtest(backtest_key, msg_data, strategy_instance_id)

        # Start the backtest with all required parameters
        self.run_backtest(
            strategy_class=MappedStrategy,
            data_feed=data_feed,
            msg_data=msg_data,
            user_name=user_name,
            socket_conn_id=socket_conn_id,
            strategy_instance=strategy_instance,
            backtest_name=backtest_name,
            user_id=user_id,
            backtest_key=backtest_key,
            strategy_name=user_strategy.get("name"),
            precomputed_indicators=precomputed_indicators
        )

        logger.info(f"Backtest '{backtest_name}' started for user '{user_name}'.")
        return {"status": "started", "backtest_name": backtest_name}

    # Define the backtest callback
    def backtest_callback(self, user_name, backtest_name, tbl_key,
                          strategy_instance_id, socket_conn_id, backtest_key, results):
        """
        Callback to handle the completion of a backtest.
        - Updates strategy stats.
        - Emits results to the client.
        - Ensures cleanup of backtest resources.
        """
        try:
            # If backtest failed
            if not results.get("success", False):
                self.store_backtest_results(backtest_key, results)
                logger.error(f"Backtest '{backtest_name}' failed for user '{user_name}': {results.get('message')}")

            # If backtest succeeded
            else:
                # Calculate additional stats
                stats = self.calculate_stats(results)
                results["stats"] = stats  # Include stats in results

                # Store results and update strategy stats
                self.store_backtest_results(backtest_key, results)
                self.update_strategy_stats(tbl_key, results)

                # Emit results to the client
                logger.debug(f"Emitting backtest results: {results}")
                sanitized_results = self.sanitize_results(results)
                self.socketio.emit(
                    'message',
                    {
                        "reply": 'backtest_results',
                        "data": {
                            'test_id': backtest_name,
                            'results': sanitized_results
                        }
                    },
                    room=socket_conn_id,
                )

                logger.info(f"Backtest '{backtest_name}' completed successfully for user '{user_name}'."
                            f"\nresults:{sanitized_results}\n")

        except Exception as e:
            logger.error(f"Error in backtest callback for '{backtest_name}': {str(e)}", exc_info=True)

        finally:
            # Ensure resources are cleaned up regardless of success or failure
            self.cleanup_backtest(backtest_key, strategy_instance_id)
            logger.info(f"Cleanup completed for backtest '{backtest_name}' (key: {backtest_key}).")

    def start_periodic_purge(self, interval_seconds: int = 3600):
        """
        Starts a background task that periodically purges expired cache entries and cleans up orphaned backtest contexts.

        :param interval_seconds: Time interval between purges in seconds.
        """
        def purge_task():
            while True:
                time.sleep(interval_seconds)
                # Purge expired entries from all caches
                for cache_name, cache in self.data_cache.caches.items():
                    if isinstance(cache, RowBasedCache) or isinstance(cache, TableBasedCache):
                        cache._purge_expired()
                        logger.debug(f"Purged expired entries from cache '{cache_name}'.")

                # Additionally, clean up orphaned backtest contexts
                self.cleanup_orphaned_backtest_contexts()

        self.socketio.start_background_task(purge_task)

    def cleanup_orphaned_backtest_contexts(self) -> None:
        """
        Identifies and removes orphaned backtest contexts that do not have corresponding entries in 'tests' cache.
        """
        try:
            # Fetch all strategy_instance_ids from 'strategy_contexts' that start with 'test_'
            strategy_contexts_df = self.data_cache.get_rows_from_datacache_with_operator(
                cache_name='strategy_contexts',
                filter_vals=[('strategy_instance_id', 'LIKE', 'test_%')]
            )

            if strategy_contexts_df.empty:
                logger.debug("No backtest contexts found for cleanup.")
                return

            # Iterate through each backtest context
            for _, row in strategy_contexts_df.iterrows():
                strategy_instance_id = row['strategy_instance_id']
                # Check if this backtest exists in 'tests' cache
                found = False
                tests_cache = self.data_cache.get_cache('tests')
                if isinstance(tests_cache, RowBasedCache):
                    for key, entry in tests_cache.cache.items():
                        if entry.data.get('strategy_instance_id') == strategy_instance_id:
                            found = True
                            break

                if not found:
                    # Orphaned context found; proceed to remove it
                    self.data_cache.remove_row_from_datacache(
                        cache_name='strategy_contexts',
                        filter_vals=[('strategy_instance_id', strategy_instance_id)],  # Correct format
                        remove_from_db=True
                    )
                    logger.info(
                        f"Orphaned backtest context '{strategy_instance_id}' removed from 'strategy_contexts' cache.")
        except Exception as e:
            logger.error(f"Error during cleanup of orphaned backtest contexts: {e}", exc_info=True)

    def shutdown_handler(self, signum, frame):
        """
        Handles server shutdown signals to perform cleanup.
        """
        logger.info(f"Received shutdown signal: {signum}. Cleaning up ongoing backtests.")
        try:
            # Fetch all backtest entries from 'tests' cache
            tests_cache = self.data_cache.get_cache('tests')
            if isinstance(tests_cache, RowBasedCache):
                active_tests = list(tests_cache.cache.keys())  # Get all backtest keys
                for backtest_key in active_tests:
                    # Fetch strategy_instance_id for each backtest
                    strategy_instance_id = self.data_cache.get_cache_item(
                        item_name='strategy_instance_id',
                        cache_name='tests',
                        filter_vals=('tbl_key', backtest_key)
                    )
                    if strategy_instance_id:
                        self.cleanup_backtest(backtest_key, strategy_instance_id)
            logger.info("Cleanup of ongoing backtests completed.")
        except Exception as e:
            logger.error(f"Error during shutdown cleanup: {e}", exc_info=True)
        finally:
            logger.info("Exiting.")
            exit(0)

    def calculate_stats(self, results: dict) -> dict:
        """
        Calculate performance metrics from backtest results.

        :param results: Dictionary containing backtest results.
        :return: Dictionary of calculated performance metrics.
        """
        # Extract required data
        initial_capital = results.get('initial_capital', 0)
        final_value = results.get('final_portfolio_value', 0)
        equity_curve = results.get('equity_curve', [])
        trades = results.get('trades', [])

        # Ensure equity curve and trades are valid
        if not equity_curve or not trades:
            return {key: 0 for key in [
                'total_return', 'sharpe_ratio', 'sortino_ratio', 'calmar_ratio',
                'volatility', 'max_drawdown', 'profit_factor', 'average_pnl',
                'number_of_trades', 'win_loss_ratio', 'max_consecutive_wins',
                'max_consecutive_losses', 'win_rate', 'loss_rate']}

        # Convert equity_curve to numpy array
        equity_curve = np.array(equity_curve, dtype=float)

        # Calculate Returns
        returns = np.array(
            [(equity_curve[i] - equity_curve[i - 1]) / equity_curve[i - 1]
             for i in range(1, len(equity_curve))],
            dtype=float
        )

        # Performance Metrics
        try:
            total_return = (final_value - initial_capital) / initial_capital * 100
        except ZeroDivisionError:
            total_return = 0

        mean_return = np.mean(returns)
        risk_free_rate = 0.0  # Adjust if needed

        # Filter downside returns and calculate downside standard deviation
        downside_returns = returns[returns < 0]
        downside_std = np.std(downside_returns) if len(downside_returns) > 0 else 1

        # Sortino and Calmar Ratios
        sortino_ratio = (mean_return - risk_free_rate) / downside_std if downside_std != 0 else 0
        max_drawdown = self.calculate_max_drawdown(equity_curve)
        calmar_ratio = total_return / abs(max_drawdown) if max_drawdown != 0 else 0

        # Volatility
        volatility = np.std(returns) * np.sqrt(252)  # Annualized

        # Trade Metrics
        total_profit = sum(trade.get('pnl', 0) for trade in trades if trade.get('pnl', 0) > 0)
        total_loss = abs(sum(trade.get('pnl', 0) for trade in trades if trade.get('pnl', 0) < 0))
        profit_factor = total_profit / total_loss if total_loss != 0 else float('inf')
        average_pnl = sum(trade.get('pnl', 0) for trade in trades) / len(trades) if trades else 0

        # Consecutive Wins and Losses
        max_consec_wins, max_consec_losses = self.calculate_max_consecutive_trades(trades)

        # Win/Loss Rates
        num_wins = sum(1 for trade in trades if trade.get('pnl', 0) > 0)
        num_losses = sum(1 for trade in trades if trade.get('pnl', 0) < 0)
        win_rate = (num_wins / len(trades)) * 100 if trades else 0
        loss_rate = (num_losses / len(trades)) * 100 if trades else 0

        # Metrics Dictionary
        return {
            'total_return': total_return,
            'sharpe_ratio': results.get('sharpe_ratio', 0),
            'sortino_ratio': sortino_ratio,
            'calmar_ratio': calmar_ratio,
            'volatility': volatility,
            'max_drawdown': max_drawdown,
            'profit_factor': profit_factor,
            'average_pnl': average_pnl,
            'number_of_trades': len(trades),
            'win_loss_ratio': num_wins / num_losses if num_losses > 0 else num_wins,
            'max_consecutive_wins': max_consec_wins,
            'max_consecutive_losses': max_consec_losses,
            'win_rate': win_rate,
            'loss_rate': loss_rate,
        }

    def sanitize_results(self, results):
        """
        Recursively sanitize a dictionary for JSON serialization.
        Converts numpy types to native Python types and handles inf/nan.
        """
        return sanitize_for_json(results)

    def sanitize_results_in_list(self, results_list):
        """Sanitize a list for JSON serialization."""
        return sanitize_for_json(results_list)

    def update_strategy_stats(self, tbl_key: str, results: dict):
        """
        Update the strategy stats with the backtest results.

        :param tbl_key: The unique identifier of the strategy.
        :param results: Dictionary containing backtest results.
        """
        strategy = self.strategies.get_strategy_by_tbl_key(tbl_key)

        if strategy:
            # Calculate metrics
            stats = self.calculate_stats(results)

            # Update the strategy's stats using the Strategies class
            self.strategies.update_stats(tbl_key, stats)

            logger.info(f"Strategy with tbl_key '{tbl_key}' stats updated successfully.")
        else:
            logger.error(f"Strategy with tbl_key '{tbl_key}' not found.")

    @staticmethod
    def calculate_max_drawdown(equity_curve) -> float:
        """
        Calculate the maximum drawdown from the equity curve.

        :param equity_curve: List or numpy array of portfolio values.
        :return: Maximum drawdown as a percentage.
        """
        # Explicitly convert to numpy array if it's not already
        equity_curve = np.array(equity_curve, dtype=float)

        # Calculate peaks and drawdowns
        peak = np.maximum.accumulate(equity_curve)
        drawdowns = (equity_curve - peak) / peak
        max_drawdown = np.min(drawdowns) * 100  # Percentage
        return max_drawdown

    @staticmethod
    def calculate_max_consecutive_trades(trades: list) -> tuple:
        """
        Calculate the maximum number of consecutive wins and losses.

        :param trades: List of trade dictionaries.
        :return: Tuple of (max_consecutive_wins, max_consecutive_losses)
        """
        max_consec_wins = 0
        max_consec_losses = 0
        current_wins = 0
        current_losses = 0

        for trade in trades:
            pnl = trade.get('pnl', 0)
            if pnl > 0:
                current_wins += 1
                current_losses = 0
            elif pnl < 0:
                current_losses += 1
                current_wins = 0
            else:
                current_wins = 0
                current_losses = 0

            max_consec_wins = max(max_consec_wins, current_wins)
            max_consec_losses = max(max_consec_losses, current_losses)

        return max_consec_wins, max_consec_losses

    def store_backtest_results(self, backtest_key: str, results: dict):
        """ Store the backtest results in the cache """
        try:
            self.data_cache.modify_cache_item(
                cache_name='tests',
                filter_vals=[('tbl_key', backtest_key)],
                field_name='results',
                new_data=str(results)  # Convert dict to string or JSON as needed
            )
            logger.info(f"Backtest results stored for key '{backtest_key}'.")
        except Exception as e:
            logger.error(f"Error storing backtest results for '{backtest_key}': {e}", exc_info=True)

    @staticmethod
    def calculate_returns(equity_curve: list) -> list:
        """
        Calculate returns based on the equity curve.
        :param equity_curve: List of portfolio values over time.
        :return: List of returns.
        """
        if not equity_curve or len(equity_curve) < 2:
            logger.warning("Insufficient data to calculate returns.")
            return []

        returns = []
        for i in range(1, len(equity_curve)):
            ret = (equity_curve[i] - equity_curve[i - 1]) / equity_curve[i - 1]
            returns.append(ret)
        logger.debug(f"Calculated returns: {returns}")
        return returns
